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

…raining
bhashini-ai · Jan 14, 2022 · 95b50d7 · 95b50d7
1 parent 6d0412f
commit 95b50d7
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/train.py b/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)