Error while training simclr model

I’m trying to implement the simclr model from the lightning ai youtube playlist but I’m getting an error when calling trainer.fit(model, train_train, val_loader).

Error message:
RuntimeError
—> 12 trainer.fit(model, train_loader, val_loader)

/usr/local/lib/python3.8/dist-packages/torch/utils/data/_utils/collate.py in collate_tensor_fn(batch, collate_fn_map)
161 storage = elem.storage().new_shared(numel, device=elem.device)
162 out = elem.new(storage).resize
(len(batch), *list(elem.size()))
→ 163 return torch.stack(batch, 0, out=out)
164
165

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

Link to colab: Google Colab