I haven’t found the memory issue yet, but for now you could try split the two stages of your training.
Basically, you would run the inference on your stage 1 models, save the scores, and keep an eye on the memory usage.
Then you would load these scores in another script and train your stage 2 model.
Since you don’t need to backpropagate through the stage 1, it shouldn’t be a problem.
Would this be an option? You could narrow the memory leak to one of the two stages.