fix cuda test criteria and dedupe similar cpu tests to OSS (#3397)

Summary:
as title, follow up addressing OSS comments, i.e. https://www.internalfb.com/diff/D67302872?dst_version_fbid=815631520895451&transaction_fbid=1694362554586457 & https://www.internalfb.com/diff/D67302872?dst_version_fbid=815631520895451&transaction_fbid=683400918136325

Pull Request resolved: #3397

Reviewed By: TroyGarden

Differential Revision: D83306596

fbshipit-source-id: 8c609c6d78b0983e2fb895bf24cf06cea413d308

Author: Fei Yu

torchrec/modules/tests/test_itep_embedding_modules.py

@@ -192,7 +192,7 @@ def generate_expected_address_lookup_buffer(
 
     # pyre-ignore[56]: Pyre was not able to infer the type of argument
     @unittest.skipIf(
-        torch.cuda.device_count() <= 1,
+        torch.cuda.device_count() < 1,
         "Not enough GPUs, this test requires at least two GPUs",
     )
     def test_init_itep_module(self) -> None:
@@ -245,7 +245,7 @@ def test_init_itep_module_without_pruned_table(self) -> None:
 
     # pyre-ignore[56]: Pyre was not able to infer the type of argument
     @unittest.skipIf(
-        torch.cuda.device_count() <= 1,
+        torch.cuda.device_count() < 1,
         "Not enough GPUs, this test requires at least one GPU",
     )
     def test_train_forward(self) -> None:
@@ -270,7 +270,7 @@ def test_train_forward(self) -> None:
 
     # pyre-ignore[56]: Pyre was not able to infer the type of argument
     @unittest.skipIf(
-        torch.cuda.device_count() <= 1,
+        torch.cuda.device_count() < 1,
         "Not enough GPUs, this test requires at least one GPU",
     )
     def test_train_forward_vbe(self) -> None:
@@ -295,7 +295,7 @@ def test_train_forward_vbe(self) -> None:
 
     # pyre-ignore[56]: Pyre was not able to infer the type of argument
     @unittest.skipIf(
-        torch.cuda.device_count() <= 1,
+        torch.cuda.device_count() < 1,
         "Not enough GPUs, this test requires at least one GPU",
     )
     # Mock out reset_weight_momentum to count calls
@@ -329,7 +329,7 @@ def test_check_pruning_schedule(
 
     # pyre-ignore[56]: Pyre was not able to infer the type of argument
     @unittest.skipIf(
-        torch.cuda.device_count() <= 1,
+        torch.cuda.device_count() < 1,
         "Not enough GPUs, this test requires at least one GPU",
     )
     # Mock out reset_weight_momentum to count calls
@@ -365,7 +365,7 @@ def test_eval_forward(
 
     # pyre-ignore[56]: Pyre was not able to infer the type of argument
     @unittest.skipIf(
-        torch.cuda.device_count() <= 1,
+        torch.cuda.device_count() < 1,
         "Not enough GPUs, this test requires at least two GPUs",
     )
     def test_iter_increment_per_forward(self) -> None:
@@ -397,7 +397,7 @@ def test_iter_increment_per_forward(self) -> None:
 
     # pyre-ignore[56]: Pyre was not able to infer the type of argument
     @unittest.skipIf(
-        torch.cuda.device_count() <= 1,
+        torch.cuda.device_count() < 1,
         "Not enough GPUs, this test requires at least one GPU",
     )
     def test_iter_passed_as_int_to_itep_module(self) -> None:
@@ -440,7 +440,7 @@ def mock_forward(features: KeyedJaggedTensor, iter_val: int) -> List[Tensor]:
 
     # pyre-ignore[56]: Pyre was not able to infer the type of argument
     @unittest.skipIf(
-        torch.cuda.device_count() <= 1,
+        torch.cuda.device_count() < 1,
         "Not enough GPUs, this test requires at least one GPU",
     )
     def test_blank_line_formatting_preserved(self) -> None:
@@ -472,7 +472,7 @@ def test_blank_line_formatting_preserved(self) -> None:
 
     # pyre-ignore[56]: Pyre was not able to infer the type of argument
     @unittest.skipIf(
-        torch.cuda.device_count() <= 1,
+        torch.cuda.device_count() < 1,
         "Not enough GPUs, this test requires at least one GPU",
     )
     def test_iter_boundary_values_with_pruning_logic(self) -> None:
@@ -523,7 +523,7 @@ def mock_forward(
         torch.cuda.device_count() <= 1,
         "Not enough GPUs, this test requires at least one GPU",
     )
-    def test_error_handling_invalid_iter_tensor_values(self) -> None:
+    def test_error_handling_invalid_iter_tensor_values_cuda(self) -> None:
         """Test behavior with invalid iter tensor values."""
         itep_module = GenericITEPModule(
             table_name_to_unpruned_hash_sizes=self._table_name_to_unpruned_hash_sizes,
@@ -567,7 +567,7 @@ def test_error_handling_invalid_iter_tensor_values(self) -> None:
 
     # pyre-ignore[56]: Pyre was not able to infer the type of argument
     @unittest.skipIf(
-        torch.cuda.device_count() <= 1,
+        torch.cuda.device_count() < 1,
         "Not enough GPUs, this test requires at least one GPU",
     )
     def test_iter_consistency_across_training_steps(self) -> None:
@@ -626,7 +626,7 @@ def track_iter_forward(
 
     # pyre-ignore[56]: Pyre was not able to infer the type of argument
     @unittest.skipIf(
-        torch.cuda.device_count() <= 1,
+        torch.cuda.device_count() < 1,
         "Not enough GPUs, this test requires at least one GPU",
     )
     def test_performance_iter_conversion_overhead(self) -> None:
@@ -680,7 +680,7 @@ def test_performance_iter_conversion_overhead(self) -> None:
 
     # pyre-ignore[56]: Pyre was not able to infer the type of argument
     @unittest.skipIf(
-        torch.cuda.device_count() <= 1,
+        torch.cuda.device_count() < 1,
         "Not enough GPUs, this test requires at least one GPU",
     )
     def test_iter_type_conversion_edge_cases(self) -> None:
@@ -744,7 +744,7 @@ def capture_iter_forward(
 
     # pyre-ignore[56]: Pyre was not able to infer the type of argument
     @unittest.skipIf(
-        torch.cuda.device_count() <= 1,
+        torch.cuda.device_count() < 1,
         "Not enough GPUs, this test requires at least one GPU",
     )
     def test_concurrent_forward_passes_iter_safety(self) -> None: