Are you seeing the OOM error after a few iterations or how were you able to see the validation step being faster?
Since Python uses function scoping, you might want to wrap some code in specific functions so that they’ll be cleared as explained here 10.