Based on the stack trace it seems your returned image tensors do not have the same shape, as the collate function raises this error.
Could you pass the image_size argument as a tuple to Resize as:
transforms.Resize((image_size, image_size))

This will make sure to return images with the specified shape in the width and height dimension.
From the docs:

Desired output size. If size is a sequence like (h, w), output size will be matched to this. If size is an int, smaller edge of the image will be matched to this number. i.e, if height > width, then image will be rescaled to (size * height / width, size)