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)
So, I have to change it to 3? The number of features or number of color channels?