RuntimeError: The size of tensor a (2) must match the size of tensor b (4) at non-singleton dimension 0

[MarkoHaralovic]

Issue: RuntimeError during Pretraining on ImageNet with Small Batch Sizes

I encountered the following error when trying to pretrain any of the models on ImageNet:

RuntimeError: The size of tensor a (2) must match the size of tensor b (4) at non-singleton dimension 0

Environment

• Setup: Single GPU environment

Issue

The issue seems to stem from the batch size configuration, specifically when the batch size is lower than 16. The model fails with the mentioned error when trying to run with small batch sizes.

Workaround

The error can be avoided by setting the batch size to any multiple of 2 that is greater than or equal to 16.

Suggested fix

It would be helpful to handle this error more gracefully or provide clear documentation that the batch size needs to be a multiple of 2 and at least 16 when using certain configurations.