Fixed wrong device type for test_psnr on hvd (pytorch#2700)

StraitRobot · Sep 5, 2022 · 34a66c5 · 34a66c5
1 parent 29d4573
commit 34a66c5
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tests/ignite/metrics/test_psnr.py b/tests/ignite/metrics/test_psnr.py
@@ -320,7 +320,8 @@ def test_distrib_xla_nprocs(xmp_executor):
 @pytest.mark.skipif(not idist.has_hvd_support, reason="Skip if no Horovod dist support")
 @pytest.mark.skipif("WORLD_SIZE" in os.environ, reason="Skip if launched as multiproc")
 def test_distrib_hvd(gloo_hvd_executor):
-    device = "cpu" if not torch.cuda.is_available() else "cuda"
+
+    device = torch.device("cpu" if not torch.cuda.is_available() else "cuda")
     nproc = 4 if not torch.cuda.is_available() else torch.cuda.device_count()
 
     gloo_hvd_executor(_test_distrib_input_float, (device,), np=nproc, do_init=True)