If you’re experiencing this OOM memory after a while, could you check, if your GPU memory usage is growing?
Also, are you storing some variables which are still attached to the computation graph?
A common mistake is to store the loss for further debugging without detaching it:
# Memory will grow
loss = criterion(output, target)
losses.append(loss)

# This should be the right way, if you don't need any ops on loss anymore
loss = criterion(output, target)
losses.append(loss.detach().item())