I used 8 GPUs and the original batch size is 128.
In debugging and reproducing this error, I used 4 GPUs and batch size of 128, and using just 278 datapoints.
When debugging, the data is split into 3 batches, the first is split like this:
`
packed torch.Size([96, 512])
packed torch.Size([96, 512])
packed torch.Size([92, 512])
packed torch.Size([96, 512])
output torch.Size([96, 512])
output torch.Size([96, 512])
output torch.Size([92, 512])
output torch.Size([96, 512]),
`
the second like this:
`
packed torch.Size([96, 512])
packed torch.Size([96, 512])
packed torch.Size([92, 512])
packed torch.Size([96, 512])
output torch.Size([96, 512])
output torch.Size([96, 512])
output torch.Size([92, 512])
output torch.Size([96, 512])
`
the third like this:
`
packed torch.Size([18, 512])
packed torch.Size([18, 512])
packed torch.Size([12, 512])
packed torch.Size([18, 512])
output torch.Size([18, 512])
output torch.Size([18, 512])
output torch.Size([12, 512])
output torch.Size([18, 512])
`
and then the error occurred when gathering after the third batch.
All the previous batches have size of 32 for each GPU (4 GPUs), the last have 6, 6, 4 and 6 batches and maybe that is the problem?