diff --git a/fp8_optimization.py b/fp8_optimization.py
index 7abe973..ae5164c 100644
--- a/fp8_optimization.py
+++ b/fp8_optimization.py
@@ -15,7 +15,7 @@ def fp8_linear_forward(cls, base_dtype, input):
             else:
                 scale_weight = scale_weight.to(input.device).squeeze()
             
-            scale_input = torch.ones((), device=input.device, dtype=torch.float32).squeeze()
+            scale_input = torch.ones((), device=input.device, dtype=torch.float32)
             
             input = torch.clamp(input, min=-448, max=448, out=input)
             inn = input.reshape(-1, input_shape[2]).to(torch.float8_e4m3fn).contiguous() #always e4m3fn because e5m2 * e5m2 is not supported