kuleshov-group
/

bd3lm-owt-block_size4

Text Generation

language-modeling

Model card Files Files and versions Community

marriola commited on 19 days ago

Commit

d5ecd4b

·

verified ·

1 Parent(s): 95c0f16

Upload BD3LM

Files changed (1) hide show

modeling_bd3lm.py +1 -1

modeling_bd3lm.py CHANGED Viewed

@@ -341,7 +341,7 @@ class DDiTBlock(nn.Module):
       qkv = self.attn_qkv(x)
     # store kv cache in a sliding window (can't exceed context len)
     if store_kv:
-      self.kv_cache = qkv[:, -(self.n-self.block_size):]
     qkv = einops.rearrange(
       qkv,

       qkv = self.attn_qkv(x)
     # store kv cache in a sliding window (can't exceed context len)
     if store_kv:
+      self.kv_cache = qkv[:, -(self.max_seqlen-self.block_size):]
     qkv = einops.rearrange(
       qkv,