Here’s a collator I use, it works for tensors of any dimension:
class ZeroPadCollator:

    @staticmethod
    def collate_tensors(batch: List[torch.Tensor]) -> torch.Tensor:
        dims = batch[0].dim()
        max_size = [max([b.size(i) for b in batch]) for i in range(dims)]
        size = (len(batch),) + tuple(max_size)
        canvas = batch[0].new_zeros(size=size)
        for i, b in enumerate(batch):
            sub_tensor = canvas[i]
            for d in range(dims):
                sub_tensor = sub_tensor.narrow(d, 0, b.size(d))
            sub_tensor.add_(b)
        return canvas

    def collate(self, batch, ) -> List[torch.Tensor]:
        dims = len(batch[0])
        return [self.collate_tensors([b[i] for b in batch]) for i in range(dims)]

Then I simply use:
    zero_pad = ZeroPadCollator()
    loader = DataLoader(train, args.batch_size, collate_fn=zero_pad.collate)```