Spaces:
Running
on
Zero
Running
on
Zero
Harry Coultas Blum
commited on
Commit
·
8b2e529
1
Parent(s):
73d9f21
trying to cast
Browse files- vui/model.py +2 -0
vui/model.py
CHANGED
|
@@ -128,6 +128,8 @@ class MHA(nn.Module):
|
|
| 128 |
k = repeat_kv(k, self.n_reps)
|
| 129 |
v = repeat_kv(v, self.n_reps)
|
| 130 |
|
|
|
|
|
|
|
| 131 |
is_causal = self.causal and self.kv_cache is None
|
| 132 |
|
| 133 |
out = F.scaled_dot_product_attention(
|
|
|
|
| 128 |
k = repeat_kv(k, self.n_reps)
|
| 129 |
v = repeat_kv(v, self.n_reps)
|
| 130 |
|
| 131 |
+
q, k, v = q.to(x.dtype), k.to(x.dtype), v.to(x.dtype)
|
| 132 |
+
|
| 133 |
is_causal = self.causal and self.kv_cache is None
|
| 134 |
|
| 135 |
out = F.scaled_dot_product_attention(
|