Based on the error message if seems that the tensors stored in batch, which should be created in the Dataset.__getitem__ method, do not have the same shape.
If you are dealing with variable input shapes, you could use a custom collate function as described here.