Update modeling_recllama.py

modeling_recllama.py

@@ -141,11 +141,12 @@ class RecDynamicCache(DynamicCache):
 
 class RecLlamaForCausalLM(LlamaForCausalLM):
     config_class = RecLlamaConfig
-    def __init__(self, config: RecLlamaConfig):
+    def __init__(self, config: RecLlamaConfig, num_steps=None):
         super().__init__(config)
         self.prelude_layers = config.prelude_layers
         self.recurrent_layers = config.recurrent_layers
         self.coda_layers = config.coda_layers
+        self.num_steps = num_steps
         
         for i in range(len(self.model.layers)):
             self.model.layers[i].self_attn.k_proj.bias = nn.Parameter(torch.randn(1, self.model.layers[i].self_attn.k_proj.out_features)) #nn.Parameter(torch.full((1, self.model.layers[i].self_attn.k_proj.out_features), k_bias_value))
@@ -374,10 +375,12 @@ class RecLlamaForCausalLM(LlamaForCausalLM):
         position_embeddings,
         num_steps=None,
     ):
-        if num_steps is None:
+        if num_steps is None and self.num_steps is None:
             num_steps_no_grad, num_steps_with_grad = self.randomized_iteration_sampler()  # type: ignore
         elif hasattr(num_steps, "__len__") and len(num_steps) > 1:
             num_steps_no_grad, num_steps_with_grad = num_steps
+        elif self.num_steps is not None:
+            num_steps_no_grad, num_steps_with_grad = self.num_steps, self.num_steps
         else:
             num_steps_no_grad, num_steps_with_grad = num_steps, torch.tensor(0)