[diffusion] optimize: preserve dtype in wanvae nearest upsample (#27077)

2026-06-03 08:32:28 +08:00
parent c55548ba11
commit 83bc776612
1 changed files with 2 additions and 0 deletions
@@ -211,6 +211,8 @@ class WanUpsample(nn.Upsample):
    """

    def forward(self, x):
+        if current_platform.is_amp_supported():
+            return super().forward(x)
        return super().forward(x.float()).type_as(x)