Update modeling_quiet.py

modeling_quiet.py

@@ -1402,17 +1402,17 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
         hidden_states_before = outputs_before[0][:, -1:, :]
 
         # two new tokens: last continuation token and end thought token
-		outputs_after = self.model(
-			input_ids=torch.cat([next_token_id.unsqueeze(-1).to(input_ids.device), torch.tensor([[end_thought_token_id]] * batch_size).to(input_ids.device)], dim=-1),
-			attention_mask=torch.cat([attention_mask[:, -1:], torch.ones((batch_size, 1)).to(attention_mask.device)], dim=-1),
-			position_ids=position_ids,
-			past_key_values=new_key_values,
-			inputs_embeds=inputs_embeds,
-			use_cache=use_cache,
-			output_attentions=output_attentions,
-			output_hidden_states=output_hidden_states,
-			return_dict=return_dict,
-		)
+        outputs_after = self.model(
+            input_ids=torch.cat([next_token_id.unsqueeze(-1).to(input_ids.device), torch.tensor([[end_thought_token_id]] * batch_size).to(input_ids.device)], dim=-1),
+            attention_mask=torch.cat([attention_mask[:, -1:], torch.ones((batch_size, 1)).to(attention_mask.device)], dim=-1),
+            position_ids=position_ids,
+            past_key_values=new_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
         hidden_states_after = outputs_after[0][:, -1:, :]
 
         # Apply the talk head to get the mixing weight