Could you explain the shape of packed and output?
It looks like the batch size is multiplied by 3 but not on all GPUs.
I understand the batches will be split in dim0, but I’m currently not sure how to debug this issue.
Any hint would be useful!