From 3167e088edca37fe2cbaa03d34dedcd53dd62c56 Mon Sep 17 00:00:00 2001 From: kijai <40791699+kijai@users.noreply.github.com> Date: Sun, 5 Oct 2025 12:26:44 +0300 Subject: [PATCH] Update fp8_optimization.py --- fp8_optimization.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/fp8_optimization.py b/fp8_optimization.py index 217ab86..7abe973 100644 --- a/fp8_optimization.py +++ b/fp8_optimization.py @@ -13,9 +13,10 @@ def fp8_linear_forward(cls, base_dtype, input): if scale_weight is None: scale_weight = torch.ones((), device=input.device, dtype=torch.float32) else: - scale_weight = scale_weight.to(input.device) + 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