I am trying to apply a different threshold to each input channel (e.g., similar to PReLU with \alpha set to nChannels), but the thresholds do not need to be learned. So far I have the following, but it is too slow due to the for loop.
class ChannelwiseThreshold(nn.Module):
    def __init__(self):
        super(ChannelwiseThreshold, self).__init__()
    
    def forward(self, x, thresholds):
        h = []
        for i in range(len(thresholds)):
            h.append(F.threshold(x[:, i], thresholds[i], 0))
        
        return torch.stack(h, dim=1)

Is there any more efficient way of achieving this?
Thanks,
Brad