Pull jik876#25 Solve imbalanced gpu memory at multi-gpu distributed training

Author: shivahr

train.py

@@ -27,6 +27,7 @@ def train(rank, a, h):
                            world_size=h.dist_config['world_size'] * h.num_gpus, rank=rank)
 
     torch.cuda.manual_seed(h.seed)
+    torch.cuda.set_device(rank)
     device = torch.device('cuda:{:d}'.format(rank))
 
     generator = Generator(h).to(device)