Image Size Changes from [640, 640] to [H, W] During Training?

[wwma]

I'm following the example from this notebook while using segmentation_models.pytorch for a semantic segmentation task.

For data augmentation, I applied common techniques like cropping and rotation. Specifically, I set the cropping size to [640, 640].

While debugging, I noticed that both the Dataset and Dataloader initially load images with the expected shape [bz, 640, 640]. However, during training, when computing the loss, the shape of image and mask changes to [bz, H, W], where H and W are the original image dimensions.

Screenshots:

I am trying to understand how this transformation happens.
Why does the image size change back from [640, 640] to the original size during training?
What could be causing this behavior?
Will this have any impact on the training process?

Any insights or suggestions would be greatly appreciated! Thanks in advance!

[qubvel]

Hi @wwma, thanks for the issue, that's indeed a strange behavior, there should be some bug somewhere, but it's hard to tell where for your specific case