Update modeling_quiet.py

modeling_quiet.py

@@ -1233,7 +1233,6 @@ class QuietForCausalLM(QuietPreTrainedModel):
         self.loss_mean = loss_mean
         self.all_rewards = []
         self.all_unreduced_losses = []
-        self.kill_after = 100
 
         self.start_embedding = nn.Parameter(torch.zeros(2, self.model.config.hidden_size))
         self.end_embedding = nn.Parameter(torch.zeros(2, self.model.config.hidden_size))
@@ -1455,6 +1454,13 @@ class QuietForCausalLM(QuietPreTrainedModel):
         >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
         "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
         ```"""
+        log_dict = self.log_dict if self.training else self.eval_log_dict
+
+        if not self.training:
+            n_ahead_talk_to_restore = self.n_ahead_talk
+            n_passes_to_restore = self.n_passes
+            self.n_ahead_talk = 1
+            self.n_passes = 1
 
         output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
         output_hidden_states = (