From bdb1d996bee26e34f59f0f7fe24f928b4bb21996 Mon Sep 17 00:00:00 2001 From: zzh <67720593+BLANKETusers@users.noreply.github.com> Date: Wed, 12 Aug 2026 10:40:05 +0800 Subject: [PATCH] [Bugfix] Fix HunyuanImage3 accuracy test (#5981) Signed-off-by: zzh <943967662@qq.com> --- vllm_omni/diffusion/layers/norm.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/vllm_omni/diffusion/layers/norm.py b/vllm_omni/diffusion/layers/norm.py index 83897eb15..8d91df8f0 100644 --- a/vllm_omni/diffusion/layers/norm.py +++ b/vllm_omni/diffusion/layers/norm.py @@ -65,9 +65,10 @@ class LayerNorm(nn.LayerNorm, CustomOp): class RMSNorm(CustomOp): - def __init__(self, hidden_size: int, eps: float = 1e-6, dtype: torch.dtype = torch.float32) -> None: + def __init__(self, hidden_size: int, eps: float = 1e-6, dtype: torch.dtype | None = None) -> None: super().__init__() - self.weight = nn.Parameter(torch.ones(hidden_size, dtype=dtype)) + weight_dtype = dtype or torch.get_default_dtype() + self.weight = nn.Parameter(torch.ones(hidden_size, dtype=weight_dtype)) self.hidden_size = hidden_size self.variance_epsilon = eps