GPU分配BUG: Duplicate GPU detected : rank 1 and rank 0 both on CUDA device d5000

#BUG

[rank1]: Last error:
[rank1]: Duplicate GPU detected : rank 1 and rank 0 both on CUDA device d5000

解决

添加

   if th.cuda.is_available():
       local_rank = comm.rank % th.cuda.device_count()
       th.cuda.set_device(local_rank)
       os.environ["LOCAL_RANK"] = str(local_rank)