minor

Signed-off-by: realAsma <[email protected]>

Author: realAsma

modelopt/torch/quantization/tensor_quant.py

@@ -80,7 +80,7 @@ def scaled_e4m3_impl(
     Returns:
         Input tensors faked quantized to FP8.
     """
-    if inputs.is_cpu or amax is None or amax.squeeze().ndim > 1:
+    if (not inputs.is_cuda) or amax is None or amax.squeeze().ndim > 1:
         return fp8_eager(inputs, amax)
 
     cuda_ext_fp8 = get_cuda_ext_fp8(raise_if_failed=False)