In that case you could lower the batch size to 1 and check, if it’s still running out of memory.
If that’s the case, keep the batch size at 1, split the batch in the DataLoader loop in dim0 and loop over smaller input tensors.
Alternatively, you could also try to create the patches from each image in several steps in the __getitem__, but I think this would be much more complicated, as you might need a custom sampler for it etc.