Hi,
If you’re designing this from scratch, I would advise to create new processes to run the inner job to ensure no issue with cleanup of remaining states.
You can check this issue for issues related to GPU OOM failures: https://github.com/pytorch/pytorch/issues/18853 6.
Also be aware that if an assert happens in a cuda kernel (like index out of range), the cuda driver cannot recover from it and you will have to restart the process. So your code example won’t help in that case as all the other trainings will fail.