Distributed training does not exit properly after completion

[Licko0909]

Dear Author,
This is a great result, I am recently using TAPE for sequence design task, I have no problem in single card training process, but I encounter a little problem using distributed training.
I'm looking forward to hearing from you.

Problem Description:

• After the distributed training is finished, the program cannot be exited normally and remains in running state
• There have been no problems with the training process.
• 4 RTX6000 on one machine

I'm looking forward to hearing from you!