Fix GPU resource check (#2659)

Summary: Pull Request resolved: #2659 Test only checks for existance of GPU, needs to check for at least 2 GPUs Reviewed By: sarckk Differential Revision: D67722949 fbshipit-source-id: 8ed02f6d4cc9fc56db0374c52ddc573f4c682401
pytorch · Dec 30, 2024 · ff03e14 · ff03e14
1 parent ab27f16
commit ff03e14
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/examples/retrieval/tests/test_two_tower_retrieval.py b/examples/retrieval/tests/test_two_tower_retrieval.py
@@ -21,8 +21,8 @@ class InferTest(unittest.TestCase):
     @skip_if_asan
     # pyre-ignore[56]
     @unittest.skipIf(
-        not torch.cuda.is_available(),
-        "this test requires a GPU",
+        torch.cuda.device_count() <= 1,
+        "Not enough GPUs, this test requires at least two GPUs",
     )
     def test_infer_function(self) -> None:
         infer(