Thank you very much ptrblck. It was really helpful.
It seems that I manage to make a function out off it for those who struggle with the issue.
def count_input_neuron(model, image_dim):
    return model(torch.rand(1, *(image_dim))).data.view(1, -1).size(1)




 ptrblck:

However, as your model expects 32 input channels, your input won’t work at all currently.


So, I have to change it to 3? The number of features or number of color channels?