Fixed the Unexpected floating ScalarType issue that occurred when used with wanvideowrapper #6

2025-12-10 02:23:08 +08:00
parent b0c9f6b131
commit dc195264eb
5 changed files with 23 additions and 7 deletions
@@ -71,8 +71,7 @@ class TransformerDecoderLayer(nn.Module):
        return tensor if pos is None else tensor + pos

    def forward_ffn(self, tgt):
-        with torch.amp.autocast(device_type="cuda", enabled=False):
-            tgt2 = self.linear2(self.dropout3(self.activation(self.linear1(tgt))))
+        tgt2 = self.linear2(self.dropout3(self.activation(self.linear1(tgt))))
        tgt = tgt + self.dropout4(tgt2)
        tgt = self.norm3(tgt)
        return tgt