diff --git a/wanvideo/modules/attention.py b/wanvideo/modules/attention.py
index 12a6cd4..ca966c8 100644
--- a/wanvideo/modules/attention.py
+++ b/wanvideo/modules/attention.py
@@ -184,9 +184,9 @@ def attention(
         #     )
         attn_mask = None
 
-        q = q.transpose(1, 2)#.to(dtype)
-        k = k.transpose(1, 2)#.to(dtype)
-        v = v.transpose(1, 2)#.to(dtype)
+        q = q.transpose(1, 2).to(dtype)
+        k = k.transpose(1, 2).to(dtype)
+        v = v.transpose(1, 2).to(dtype)
 
         out = torch.nn.functional.scaled_dot_product_attention(
             q, k, v, attn_mask=attn_mask, is_causal=causal, dropout_p=dropout_p)
@@ -196,9 +196,9 @@ def attention(
     elif attention_mode == 'sageattn':
         attn_mask = None
 
-        q = q.transpose(1, 2)#.to(dtype)
-        k = k.transpose(1, 2)#.to(dtype)
-        v = v.transpose(1, 2)#.to(dtype)
+        q = q.transpose(1, 2).to(dtype)
+        k = k.transpose(1, 2).to(dtype)
+        v = v.transpose(1, 2).to(dtype)
 
         out = sageattn_func(
             q, k, v, attn_mask=attn_mask, is_causal=causal, dropout_p=dropout_p)