The example code is:
# input_size: N * 1024 * 1
output = torch.squeeze(input)
# output_size: N * 1024

where N is batch_size.

When the code is running with single GPU mode, the output size is correct, i.e., N * 1024.
However, when using multiple GPUs (for example, 4 GPUs), the output size is wired:

When N is large, e.g. 64, the output size is correct.
When N is small, e.g. N=2, the output size will be 2048…



PS, when using output = input.squeeze(-1), the output is always correct.