test updates per review

Mike Ruberry · Mike Ruberry · commit 8419dafda9b7 · 2020-07-24T21:14:13.000-07:00
diff --git a/test/test_sparse.py b/test/test_sparse.py
@@ -1973,35 +1973,28 @@ def test_factory_type_inference(self):
     @cuda_only
     def test_factory_device_type_inference(self):
         # both indices/values are CUDA
-        shape = (1, 3)
-        for indices_device in ['cuda', 'cpu']:
-            for values_device in ['cuda', 'cpu']:
-                for sparse_device in ['cuda', 'cpu', None]:
-                    for test_empty_tensor in [True, False]:
-                        if test_empty_tensor:
-                            if indices_device != values_device and sparse_device is None:
-                                with self.assertRaises(RuntimeError):
-                                    torch.sparse_coo_tensor(torch.tensor(([0], [2]), device=indices_device),
-                                                            self.value_empty(1, 0).to(values_device),
-                                                            (1, 3, 0), device=sparse_device)
-                                continue
-                            else:
-                                t = torch.sparse_coo_tensor(torch.tensor(([0], [2]), device=indices_device),
-                                                            self.value_empty(1, 0).to(values_device),
-                                                            (1, 3, 0), device=sparse_device)
-                        else:
-                            if indices_device != values_device and sparse_device is None:
-                                with self.assertRaises(RuntimeError):
-                                    torch.sparse_coo_tensor(torch.tensor(([0], [2]), device=indices_device),
-                                                            self.value_empty(1, 0).to(values_device),
-                                                            (1, 3, 0), device=sparse_device)
-                                continue
-                            else:
-                                t = torch.sparse_coo_tensor(torch.tensor(([0], [2]), device=indices_device),
-                                                            torch.tensor([1.], device=values_device),
-                                                            (1, 3), device=sparse_device)
-                        should_be_cuda = sparse_device == 'cuda' or (sparse_device is None and values_device == 'cuda')
-                        self.assertEqual(should_be_cuda, t.is_cuda)
+
+        cpu_cuda = ('cpu', 'cuda')
+        cpu_cuda_none = cpu_cuda + (None,)
+        for indices_device, values_device, device in itertools.product(cpu_cuda,
+                                                                       cpu_cuda,
+                                                                       cpu_cuda_none):
+            indices = torch.tensor(([0], [2]), device=indices_device)
+            values = torch.tensor([1.], device=values_device)
+            empty_values = self.value_empty(1, 0).to(values_device)
+            shape = (1, 3)
+            empty_shape = (1, 3, 0)
+            if device is None and indices_device != values_device:
+                with self.assertRaises(RuntimeError):
+                    torch.sparse_coo_tensor(indices, values, shape, device=device)
+                with self.assertRaises(RuntimeError):
+                    torch.sparse_coo_tensor(indices, empty_values, empty_shape, device=device)
+            else:
+                t = torch.sparse_coo_tensor(indices, values, shape, device=device)
+                t_empty = torch.sparse_coo_tensor(indices, empty_values, empty_shape, device=device)
+                should_be_cuda = (device == 'cuda' or (device is None and values_device == 'cuda'))
+                self.assertEqual(should_be_cuda, t.is_cuda)
+                self.assertEqual(t.is_cuda, t_empty.is_cuda)
 
     @cpu_only
     def test_factory_copy(self):
diff --git a/test/test_torch.py b/test/test_torch.py
@@ -5502,15 +5502,17 @@ def test_tensor_ctor_device_inference(self, device):
         torch_device = torch.device(device)
         values = torch.tensor((1, 2, 3), device=device)
 
-        # Tests tensor (suppresses warnings)
-        with warnings.catch_warnings():
-            warnings.simplefilter("ignore")
-            self.assertEqual(torch.tensor(values).device, torch_device)
-            self.assertEqual(torch.tensor(values, dtype=torch.float64).device, torch_device)
+        # Tests tensor and as_tensor
+        # Note: warnings are suppressed (suppresses warnings)
+        for op in (torch.tensor, torch.as_tensor):
+            with warnings.catch_warnings():
+                warnings.simplefilter("ignore")
+                self.assertEqual(op(values).device, torch_device)
+                self.assertEqual(op(values, dtype=torch.float64).device, torch_device)
 
-        # Tests as_tensor
-        self.assertEqual(torch.as_tensor(values).device, torch_device)
-        self.assertEqual(torch.as_tensor(values, dtype=torch.float64).device, torch_device)
+                if self.device_type == 'cuda':
+                    with torch.cuda.device(device):
+                        self.assertEqual(op(values.cpu()).device, torch.device('cpu'))
 
         # Tests sparse ctor
         indices = torch.tensor([[0, 1, 1],
@@ -5524,6 +5526,12 @@ def test_tensor_ctor_device_inference(self, device):
         sparse_with_dtype = torch.sparse_coo_tensor(indices, values, sparse_size, dtype=torch.float64)
         self.assertEqual(sparse_with_dtype.device, torch_device)
 
+        if self.device_type == 'cuda':
+            with torch.cuda.device(device):
+                sparse_with_dtype = torch.sparse_coo_tensor(indices.cpu(), values.cpu(),
+                                                            sparse_size, dtype=torch.float64)
+                self.assertEqual(sparse_with_dtype.device, torch.device('cpu'))
+
     @unittest.skipIf(not TEST_NUMPY, "NumPy not found")
     @dtypes(torch.complex64, torch.complex128)
     def test_abs_angle_complex_to_float(self, device, dtype):
