Solved by ptrblck in post #2 
              
              
                torch.squeeze will squeeze all dimensions with size 1: 
x = torch.randn(1, 2, 1, 2, 1, 2, 1)
print(x.squeeze().shape)
> torch.Size([2, 2, 2])

In the case of a small batch size, each device might get a chunk where N==1, which will also squeeze the batch dimension. 
To avoid such errors, I would reco…