Update modeling_tiny_mixtral.py

modeling_tiny_mixtral.py

@@ -307,6 +307,10 @@ class SimpleMultiHeadAttention(nn.Module):
             
             q = q_rotary.transpose(1, 2)  # Back to (batch_size, num_heads, seq_len, head_dim)
             k = k_rotary.transpose(1, 2)  # Back to (batch_size, num_heads, seq_len, head_dim)
+            
+            q = q.to(self.device)
+            k = k.to(self.device)
+            v = v.to(self.device)
         
         # Compute attention
         if self.flash: