From 05ac1407ed890d61ac16c11dbfdbc8bb07d4586f Mon Sep 17 00:00:00 2001 From: Jongchan Park Date: Fri, 12 Oct 2018 14:54:02 +0900 Subject: [PATCH] Update main_horovod.py --- examples/cifar100/main_horovod.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/cifar100/main_horovod.py b/examples/cifar100/main_horovod.py index ef9dda7..92ff704 100644 --- 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