[Bugfix] Fix HunyuanImage3 accuracy test (#5981)
Signed-off-by: zzh <943967662@qq.com>
This commit is contained in:
1 file changed
+3
-2
@@ -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
|
||||
|
||||
|
||||
Reference in new issue
Block a user