kuleshov-group
/

bd3lm-owt-block_size16

Text Generation

language-modeling

Model card Files Files and versions

marriola commited on Mar 4

Commit

e58f74c

·

verified ·

1 Parent(s): d25f94c

Upload BD3LM

Files changed (1) hide show

modeling_bd3lm.py +4 -0

modeling_bd3lm.py CHANGED Viewed

@@ -525,6 +525,10 @@ class BD3LM(transformers.PreTrainedModel):
         'sampling_eps_max',
         torch.tensor(config.sampling_eps_max))
   def forward(
       self,
       input_ids: torch.LongTensor = None,

         'sampling_eps_max',
         torch.tensor(config.sampling_eps_max))
+  def reset_kv_cache(self):
+    for block in self.backbone.blocks:
+      block.kv_cache = None
   def forward(
       self,
       input_ids: torch.LongTensor = None,