The new code snippet isn’t consistent with the previous one, as you were previously using an int and not a mask.
        self.drop = DropFeatures(input_dim)
        self.fc1 = nn.Linear(input_dim, 1)

This indicates that input_dim is an int, not a tensor.



 Marco_Repetto:

so in this configuration, the mask is applied only on evaluation?


No, since you are not checking the self.training attribute to switch between training and evaluation mode.



 Marco_Repetto:

The problem is that if I train the model setting new_mask = torch.zeros(1, 69) the loss is actually decreasing.


I don’t know where you are setting this, but yes, if you are using a mask with all-zeros, the output will also be all zero.