Update main_horovod.py

Jongchan · Oct 12, 2018 · 05ac140 · 05ac140
1 parent 6ceea4c
commit 05ac140
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/examples/cifar100/main_horovod.py b/examples/cifar100/main_horovod.py
@@ -124,7 +124,7 @@ def forward(self, x):
 
 print ("use cuda!!")
 print ("local rank {}, rank {}".format(hvd.local_rank(),hvd.rank()))
-torch.cuda.set_device(hvd.rank())
+torch.cuda.set_device(hvd.local_rank())
 torch.cuda.manual_seed(1111)
 
 best_acc = 0