Is the OOM happening at the very first batch of the very first epoch or is it happening after a while?
There shouldn’t be any issues with GPU caching provided there isn’t some sort of unintentional “leak” somewhere due to tensors being kept around when they aren’t used anymore. From the dataloader usage here only a few batches at most should be on the GPU.
You might want to remove parts of the training loop to incrementally check if any functions are causing the memory usage to grow with every iteration or if the memory usage is a lot higher than expected for a specific function.