Can you make it simpler than this?

I saw a paper and implemented a similar coding. Can it be made simpler than this? For example, not use the for loop. The model looks like this, and this is what I made.

image

class FeatureSelectionNetwork(nn.Module):
    def __init__(self, input_size):
        super().__init__()
        
        self.input_size = input_size
        self.params = nn.ParameterList([nn.Parameter(torch.randn(())) for i in range(input_size)])
            
            
    def forward(self, x):
        if self.input_size != x.shape[1]:
            raise Exception('input size must be same')
        
        wx = []
        for i, param in enumerate(self.params):
            wx.append(x[:, i] * param)
        
        return F.relu(torch.stack(wx, dim=1))

image

could you make it faster? And do you think this layer can do feature selection in neural networks?