Update auto_fp8/quantize.py

neuralmagic · May 23, 2024 · 2c70d7a · 2c70d7a
1 parent 249902a
commit 2c70d7a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/auto_fp8/quantize.py b/auto_fp8/quantize.py
@@ -81,7 +81,7 @@ def fp8_gemm(A, A_scale, B, B_scale, bias, out_dtype):
             bias=bias,
         )
         if need_reshape:
-            output = output.reshape((batch_size, *output.shape))
+            output = output.reshape((batch_size, output.shape[0] // batch_size, output.shape[1]))
     else:
         output = torch.nn.functional.linear(
             A.to(out_dtype) * A_scale,