skip unit test if no cuda

Summary: Test Plan: Reviewers: Subscribers: Tasks: Tags:
weifengpy · Apr 28, 2024 · 6f834ce · 6f834ce
1 parent 8926ee1
commit 6f834ce
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/test/dtypes/test_nf4.py b/test/dtypes/test_nf4.py
@@ -371,16 +371,17 @@ def test_tensor_as_strided_invalid(self, input_size: Union[Tuple[int], int]):
                 stride = (nf4_tensor.stride()[1], nf4_tensor.stride()[0])
                 torch.as_strided(nf4_tensor, nf4_tensor.size(), stride, nf4_tensor.storage_offset())
 
+    @unittest.skipIf(not torch.cuda.is_available(), "Need CUDA available")
     def test_pin_memory(self):
         nf4_tensor = to_nf4(torch.randn(512 * 512))
         self.assertFalse(nf4_tensor.is_pinned())
 
         nf4_tensor = nf4_tensor.pin_memory()
         self.assertTrue(nf4_tensor.is_pinned())
 
-        if torch.cuda.is_available():
-            nf4_tensor = to_nf4(torch.randn(512 * 512, device='cuda'))
-            self.assertFalse(nf4_tensor.is_pinned())
+        nf4_tensor = to_nf4(torch.randn(512 * 512, device='cuda'))
+        self.assertFalse(nf4_tensor.is_pinned())
+
 
     @unittest.skipIf(not torch.cuda.is_available(), "Need CUDA available")
     def test_to_cuda(self):