rank instead of local rank

Jongchan · Oct 10, 2018 · 6ceea4c · 6ceea4c
1 parent cc5afbe
commit 6ceea4c
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/examples/cifar100/main_horovod.py b/examples/cifar100/main_horovod.py
@@ -123,8 +123,8 @@ def forward(self, x):
 hvd.init()
 
 print ("use cuda!!")
-print ("local rank {}".format(hvd.local_rank()))
-torch.cuda.set_device(hvd.local_rank())
+print ("local rank {}, rank {}".format(hvd.local_rank(),hvd.rank()))
+torch.cuda.set_device(hvd.rank())
 torch.cuda.manual_seed(1111)
 
 best_acc = 0