Upload BD3LM

modeling_bd3lm.py

@@ -599,10 +599,10 @@ class BD3LM(transformers.PreTrainedModel):
         'sampling_eps_max',
         torch.tensor(config.sampling_eps_max))
     
-  def reset_kv_cache(self):
+  def reset_kv_cache(self, eval_batch_size=1):
     for block in self.backbone.blocks:
       block.kv_cache = torch.zeros(
-        self.config.loader.eval_batch_size,
+        eval_batch_size,
         self.n,
         self.config.model.hidden_size * 3,
         device='cuda',