Avi:

the bug is hiding in the BN layer, which couldn’t get 1*1 shape tensor in “train mode”. but the first AdaptiveAvgPool2d layer is forcing that shape.


I don’t think this is necessarily an error, but might be a design decision.
E.g. the inception_v3 model also doesn’t accept single samples in train() mode for the same reason and fails in a batchnorm layer as seen here:
# setup
model = models.inception_v3(init_weights=False)

x = torch.randn(1, 3, 299, 299)
out = model(x) # error in batchnorm layer
> ValueError: Expected more than 1 value per channel when training, got input size torch.Size([1, 768, 1, 1])

x = torch.randn(2, 3, 299, 299)
out = model(x) # works