another variation to placate typedef int for device_count function

facebookresearch · Oct 6, 2023 · 50e521a · 50e521a
1 parent e9dab64
commit 50e521a
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/tests/test_utils.py b/tests/test_utils.py
@@ -23,9 +23,8 @@ def gpu_test(gpu_count: int = 1):
     required amount of GPU is not available
     """
     message = f"Not enough GPUs to run the test: requires {gpu_count}"
-    return pytest.mark.skipif(
-        int(torch.cuda.device_count()) < gpu_count, reason=message
-    )
+    local_gpu_count: int = torch.cuda.device_count()
+    return pytest.mark.skipif(local_gpu_count < gpu_count, reason=message)
 
 
 def init_distributed_on_file(world_size: int, gpu_id: int, sync_file: str):