wassname/iris_bigvae
1
0
Fork 0
mirror of https://github.com/wassname/iris_bigvae.git synced 2026-09-09 11:24:31 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix kv_cache faking?

Browse Source
This commit is contained in:
wassname
2023-11-13 08:11:13 +08:00
parent b97a499a19
commit e417790413
1 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/models/transformer.py
+4 -4
View File
@@ -130,10 +130,10 @@ class Transformer(nn.Module):
# fake it, since it's used to keep track of steps
if past_keys_values is not None:
k_size = past_keys_values[0]._k_cache._cache.size()
k_size = (*k_size[:2], 1, *k_size[3:])
v_size = past_keys_values[0]._v_cache._cache.size()
v_size = (*v_size[:2], 1, *v_size[3:])
past_keys_values[0].update(torch.rand(k_size), torch.rand(v_size))
# k_size = (x.shape[0], x.shape[1], x.shape[1], 1)
# v_size = past_keys_values[0]._v_cache._cache.size()
v_size = (k_size[0], k_size[1], x.shape[1], k_size[3])
past_keys_values[0].update(torch.rand(v_size), torch.rand(v_size))
return x