RuntimeError: stack expects each tensor to be equal size, but got [3, 224, 224] at entry 0 and [3, 224, 336] at entry 3

Note that Resize will behave differently on input images with a different height and width.
From the docs:

size ( sequence or int ) – 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)

If you are dealing with such images, pass the size argument as a tuple:

transforms.Resize((img_size, img_size))

CC @pr6dA