Fixing typing hint

Xilinx · Oct 28, 2024 · 4fadb9d · 4fadb9d
1 parent 1bfb332
commit 4fadb9d
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/brevitas/quant_tensor/test_quant_tensor.py b/tests/brevitas/quant_tensor/test_quant_tensor.py
@@ -55,7 +55,7 @@ def to_mx_quant_tensor(
         exponent_bit_width=4,
         mantissa_bit_width=3,
         group_size=32,
-        group_dim=1) -> FloatQuantTensor:
+        group_dim=1) -> GroupwiseFloatQuantTensor:
     mod = QuantIdentity(
         bit_width=bit_width,
         group_size=group_size,