sharif_sagar:

When i try to run the code above for other than 64*64 image size …I get the following error:


My best guess is that you are using a view or reshape operation inside the model’s forward, which might be using hard-coded and wrong shapes.
Based on the error message it seems that you are changing the batch size, so I guess in your code a similar operation such as:
x = x.view(-1, 2)

might be used while:
x = x.view(x.size(0), -1)

would be the recommended way.
In any case, this change might just yield another shape mismatch if you haven’t written the model in a way it can accept variable spatial input shapes. Could you describe a bit more how the model is handling these different spatial shapes (e.g. via adaptive pooling layers)?