From 5f3843042300ce64470e8f5876e82c7c4599797a Mon Sep 17 00:00:00 2001 From: biondizzle Date: Mon, 1 Jun 2026 11:05:09 +0000 Subject: [PATCH] Fix: use 1-dim tensors for gate_ws2 and gate_input_scale --- single_shot_inference.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/single_shot_inference.py b/single_shot_inference.py index 52fcbe13..fc19ce4d 100644 --- a/single_shot_inference.py +++ b/single_shot_inference.py @@ -723,8 +723,8 @@ def main(): router.load_nvfp4_fused_gate( gate_weight=gw_fp4, gate_weight_scale=gw_sf, - gate_ws2=torch.tensor(gw_gs, device=dev), # gsb = weight global scale - gate_input_scale=torch.tensor(1.0 / (6.0 * 448.0), device=dev), # gsa = activation global scale + gate_ws2=torch.tensor([gw_gs], device=dev, dtype=torch.float32), # gsb = weight global scale + gate_input_scale=torch.tensor([1.0 / (6.0 * 448.0)], device=dev, dtype=torch.float32), # gsa = activation global scale ) else: router.load_weights(e_bias=eb.to(dev, torch.float32))