Kenzd:

torch.cuda.empty_cache() after every batch, but it doesn’t work out.


This would slow down your code as synchronizing calls are executed to free and reallocate the memory.
Are you seeing an increasing memory usage on the GPU(s)? If so, you might accidentally store the computation graph, which would increase the memory and runtime of your training, e.g. via total_loss += loss. If you store the loss or any other tensor, which is attached to the computation graph, you should .detach() it before.