From 9d9f98cb728c729cf3c20f4fd138997c315558ce Mon Sep 17 00:00:00 2001
From: lspindler <lspindler@nvidia.com>
Date: Tue, 28 Oct 2025 08:02:26 +0100
Subject: [PATCH] Fix unittests for CPU build

---
 tests-unit/comfy_quant/test_mixed_precision.py |  7 +++++++
 tests-unit/comfy_quant/test_quant_registry.py  | 12 ++++++++++--
 2 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/tests-unit/comfy_quant/test_mixed_precision.py b/tests-unit/comfy_quant/test_mixed_precision.py
index 1102f9bd4..267bc177b 100644
--- a/tests-unit/comfy_quant/test_mixed_precision.py
+++ b/tests-unit/comfy_quant/test_mixed_precision.py
@@ -6,6 +6,13 @@ import os
 # Add comfy to path
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))
 
+def has_gpu():
+    return torch.cuda.is_available()
+
+from comfy.cli_args import args
+if not has_gpu():
+    args.cpu = True
+
 from comfy import ops
 from comfy.quant_ops import QuantizedTensor, TensorCoreFP8Layout
 
diff --git a/tests-unit/comfy_quant/test_quant_registry.py b/tests-unit/comfy_quant/test_quant_registry.py
index 2d7d3fa28..477811029 100644
--- a/tests-unit/comfy_quant/test_quant_registry.py
+++ b/tests-unit/comfy_quant/test_quant_registry.py
@@ -6,6 +6,13 @@ import os
 # Add comfy to path
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))
 
+def has_gpu():
+    return torch.cuda.is_available()
+
+from comfy.cli_args import args
+if not has_gpu():
+    args.cpu = True
+
 from comfy.quant_ops import QuantizedTensor, TensorCoreFP8Layout
 
 
@@ -49,7 +56,7 @@ class TestQuantizedTensor(unittest.TestCase):
             float_tensor,
             TensorCoreFP8Layout,
             scale=scale,
-            fp8_dtype=torch.float8_e4m3fn
+            dtype=torch.float8_e4m3fn
         )
 
         self.assertIsInstance(qt, QuantizedTensor)
@@ -96,6 +103,7 @@ class TestGenericUtilities(unittest.TestCase):
         # Verify it's a deep copy
         self.assertIsNot(qt_cloned._qdata, qt._qdata)
 
+    @unittest.skipUnless(has_gpu(), "GPU not available")
     def test_to_device(self):
         """Test device transfer"""
         fp8_data = torch.randn(10, 20, dtype=torch.float32).to(torch.float8_e4m3fn)
@@ -163,7 +171,7 @@ class TestFallbackMechanism(unittest.TestCase):
             a_fp32,
             TensorCoreFP8Layout,
             scale=scale,
-            fp8_dtype=torch.float8_e4m3fn
+            dtype=torch.float8_e4m3fn
         )
 
         # Call an operation that doesn't have a registered handler