Fix pre/postscale test for PyTorch HOROVOD_MIXED_INSTALL case.

Signed-off-by: Josh Romero <joshr@nvidia.com>
horovod · Aug 10, 2020 · a27ff94 · a27ff94
1 parent f108f86
commit a27ff94
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/test/test_torch.py b/test/test_torch.py
@@ -348,8 +348,8 @@ def test_horovod_allreduce_prescale(self):
                                    prescale_factor=factor)
 
             factor = torch.tensor(factor, dtype=torch.float64)
+            factor = factor.cuda(hvd.local_rank()) if dtype.is_cuda else factor
             if dtype.is_cuda and not int(os.environ.get('HOROVOD_MIXED_INSTALL', 0)):
-              factor = factor.cuda(hvd.local_rank())
               # For integer types, scaling done in FP64
               factor = factor.type(torch.float64 if dtype in int_types else dtype)
               tensor = tensor.type(torch.float64 if dtype in int_types else dtype)
@@ -402,8 +402,8 @@ def test_horovod_allreduce_postscale(self):
                                    postscale_factor=factor)
 
             factor = torch.tensor(factor, dtype=torch.float64)
+            factor = factor.cuda(hvd.local_rank()) if dtype.is_cuda else factor
             if dtype.is_cuda and not int(os.environ.get('HOROVOD_MIXED_INSTALL', 0)):
-              factor = factor.cuda(hvd.local_rank())
               # For integer types, scaling done in FP64
               factor = factor.type(torch.float64 if dtype in int_types else dtype)
               tensor = tensor.type(torch.float64 if dtype in int_types else dtype)