fix: dtype

blocks_jvlm.py

@@ -944,9 +944,9 @@ class MHSDPA(nn.Module):
             q, k, v = qkv.split(self.fused_dims, dim=-1)
         else:
             assert xk is not None
-            q = f.linear(xq, self.q_w.weight, self.q_b)
+            q = f.linear(xq.to(self.q_w.weight.dtype), self.q_w.weight, self.q_b)
             kv_b = torch.cat((self.k_b, self.v_b))
-            kv = f.linear(xk, self.kv_w.weight, kv_b)
+            kv = f.linear(xk.to(self.kv_w.weight.dtype), self.kv_w.weight, kv_b)
             if self.clip_qkv is not None:
                 q.clamp_(min=-self.clip_qkv, max=self.clip_qkv)
                 kv.clamp_(min=-self.clip_qkv, max=self.clip_qkv)