minor

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

Author: realAsma

modelopt/torch/quantization/tensor_quant.py

@@ -80,11 +80,11 @@ def scaled_e4m3_impl(
     Returns:
         Input tensors faked quantized to FP8.
     """
-    if inputs.is_cpu:
+    if inputs.is_cpu 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)
-    if cuda_ext_fp8 is None or amax is None:
+    if cuda_ext_fp8 is None:
         return fp8_eager(inputs, amax)
 
     with torch.cuda.device(
@@ -95,9 +95,7 @@ def scaled_e4m3_impl(
         elif amax.squeeze().ndim == 1:
             axis = amax.shape.index(amax.numel())
             outputs = cuda_ext_fp8.fake_e4m3fy_with_axis(inputs, amax.squeeze(), axis)
-        else:
-            outputs = fp8_eager(inputs, amax)
-        return outputs
+    return outputs
 
 
 def fake_quant_impl(