Harry Coultas Blum
commited on
Commit
·
dbb5fad
1
Parent(s):
17ad867
Trying to fix kvcache
Browse files- vui/model.py +3 -2
vui/model.py
CHANGED
|
@@ -39,8 +39,9 @@ class KVCache(nn.Module):
|
|
| 39 |
k_out = self.k_cache
|
| 40 |
v_out = self.v_cache
|
| 41 |
input_pos = input_pos.int()
|
| 42 |
-
k_out[:, :, input_pos] = k_val
|
| 43 |
-
v_out[:, :, input_pos] = v_val
|
|
|
|
| 44 |
|
| 45 |
return k_out, v_out
|
| 46 |
|
|
|
|
| 39 |
k_out = self.k_cache
|
| 40 |
v_out = self.v_cache
|
| 41 |
input_pos = input_pos.int()
|
| 42 |
+
k_out[:, :, input_pos] = k_val.to(k_out.dtype)
|
| 43 |
+
v_out[:, :, input_pos] = v_val.to(k_out.dtype)
|
| 44 |
+
|
| 45 |
|
| 46 |
return k_out, v_out
|
| 47 |
|