Spaces:
Running
on
Zero
Running
on
Zero
Harry Coultas Blum
commited on
Commit
·
6b554a6
1
Parent(s):
9d09b10
Trying to fix kvcache
Browse files- requirements.txt +2 -2
- vui/model.py +2 -2
requirements.txt
CHANGED
|
@@ -12,7 +12,7 @@ pyannote.audio
|
|
| 12 |
soundfile
|
| 13 |
sphn
|
| 14 |
tiktoken
|
| 15 |
-
torch
|
| 16 |
-
torchaudio
|
| 17 |
tqdm
|
| 18 |
transformers
|
|
|
|
| 12 |
soundfile
|
| 13 |
sphn
|
| 14 |
tiktoken
|
| 15 |
+
torch --index-url https://download.pytorch.org/whl/cu128
|
| 16 |
+
torchaudio --index-url https://download.pytorch.org/whl/cu128
|
| 17 |
tqdm
|
| 18 |
transformers
|
vui/model.py
CHANGED
|
@@ -38,8 +38,8 @@ class KVCache(nn.Module):
|
|
| 38 |
|
| 39 |
k_out = self.k_cache
|
| 40 |
v_out = self.v_cache
|
| 41 |
-
k_out[:, :, input_pos] = k_val
|
| 42 |
-
v_out[:, :, input_pos] = v_val
|
| 43 |
|
| 44 |
return k_out, v_out
|
| 45 |
|
|
|
|
| 38 |
|
| 39 |
k_out = self.k_cache
|
| 40 |
v_out = self.v_cache
|
| 41 |
+
k_out[:, :, input_pos] = k_val.to(k_out.dtype)
|
| 42 |
+
v_out[:, :, input_pos] = v_val.to(v_val.dtype)
|
| 43 |
|
| 44 |
return k_out, v_out
|
| 45 |
|