Fix unit test nit

microsoft · Jan 7, 2025 · 215eb63 · 215eb63
1 parent 5eb8f0b
commit 215eb63
Showing 1 changed file with 14 additions and 15 deletions.
diff --git a/test/unit_test/model/test_hf_config.py b/test/unit_test/model/test_hf_config.py
@@ -41,23 +41,22 @@ def test_device_map(self, inputs, inner):
         assert args.device_map == {"": inner}
 
     @pytest.mark.skipif(not torch.cuda.is_available(), reason="No GPU available")
-    class TestHfLoadKwargsGPU:
-        @pytest.mark.parametrize(
-            ("inputs", "inner"),
-            [
-                ("cuda:0", "cuda:0"),
-                ("0", "cuda:0"),
-            ],
-        )
-        def test_device_map_cpu(self, inputs, inner):
-            if inputs == "0":
-                inputs = torch.device(0)
+    @pytest.mark.parametrize(
+        ("inputs", "inner"),
+        [
+            ("cuda:0", "cuda:0"),
+            ("0", "cuda:0"),
+        ],
+    )
+    def test_device_map_cpu(self, inputs, inner):
+        if inputs == "0":
+            inputs = torch.device(0)
 
-            args = HfLoadKwargs(device_map=inputs)
-            assert args.device_map == inner
+        args = HfLoadKwargs(device_map=inputs)
+        assert args.device_map == inner
 
-            args = HfLoadKwargs(device_map={"": inputs})
-            assert args.device_map == {"": inner}
+        args = HfLoadKwargs(device_map={"": inputs})
+        assert args.device_map == {"": inner}
 
     @pytest.mark.parametrize(
         ("quantization_method", "quantization_config", "valid"),