Have you tried reducing the batch size? I’d suggest doing that and then monitoring the GPU memory usage levels just to see if it still happens at the second iteration and if there’s a weird pattern. I haven’t seen this with pytorch, just trying to spur some ideas