Solved by ptrblck in post #2 
              
              
                I don’t think it would make a difference, but you could pass the expected input tensor to the __init__ method and get the feature dimension via input_dim = x.size(1) assuming the features are in dim1. 
Alternatively, you could also reuse your code and use model = MLP(x.size(1), ...) instead.