c_cj:

but it causes deadlock in torch DDP after I save the model.


Does this mean that you are able to successfully save the model, but the next iteration hangs?
Could you add print statements to the code to check where exactly the script is hanging?