[diffusion] optimize: preserve dtype in wanvae nearest upsample (#27077)
This commit is contained in:
@@ -211,6 +211,8 @@ class WanUpsample(nn.Upsample):
|
|||||||
"""
|
"""
|
||||||
|
|
||||||
def forward(self, x):
|
def forward(self, x):
|
||||||
|
if current_platform.is_amp_supported():
|
||||||
|
return super().forward(x)
|
||||||
return super().forward(x.float()).type_as(x)
|
return super().forward(x.float()).type_as(x)
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user