Spaces:
Runtime error
Runtime error
Update wan/modules/causal_model.py
Browse files
wan/modules/causal_model.py
CHANGED
|
@@ -236,6 +236,7 @@ class CausalWanSelfAttention(nn.Module):
|
|
| 236 |
|
| 237 |
# output
|
| 238 |
x = x.flatten(2)
|
|
|
|
| 239 |
x = self.o(x)
|
| 240 |
return x
|
| 241 |
|
|
|
|
| 236 |
|
| 237 |
# output
|
| 238 |
x = x.flatten(2)
|
| 239 |
+
x = x.to(self.o.weight.dtype)
|
| 240 |
x = self.o(x)
|
| 241 |
return x
|
| 242 |
|