Spaces:
Running
on
Zero
Running
on
Zero
Harry Coultas Blum
commited on
Commit
·
76a4a55
1
Parent(s):
8760fa5
wip
Browse files- vui/model.py +3 -2
vui/model.py
CHANGED
@@ -38,8 +38,9 @@ class KVCache(nn.Module):
|
|
38 |
|
39 |
k_out = self.k_cache
|
40 |
v_out = self.v_cache
|
41 |
-
|
42 |
-
|
|
|
43 |
|
44 |
return k_out, v_out
|
45 |
|
|
|
38 |
|
39 |
k_out = self.k_cache
|
40 |
v_out = self.v_cache
|
41 |
+
input_pos = input_pos.to(k_out.dtype)
|
42 |
+
k_out[:, :, input_pos] = k_val
|
43 |
+
v_out[:, :, input_pos] = v_val
|
44 |
|
45 |
return k_out, v_out
|
46 |
|