Update modeling_quiet.py

modeling_quiet.py

@@ -1212,11 +1212,11 @@ class QuietForCausalLM(QuietPreTrainedModel):
             use_cache=use_cache,
             output_attentions=output_attentions,
             output_hidden_states=output_hidden_states,
-            return_dict=return_dict,
+            return_dict=True,  # Set return_dict=True
         )
 
         hidden_states = outputs.last_hidden_state
-        base_logits = outputs.logits
+        logits = self.lm_head(hidden_states)
 
         thought_ids, thought_embeddings = self.model._generate_thoughts(hidden_states, max_length=self.thought_length)
         thought_hidden_states = self.model(inputs_embeds=thought_embeddings).last_hidden_state
@@ -1249,7 +1249,7 @@ class QuietForCausalLM(QuietPreTrainedModel):
 
         return CausalLMOutputWithPast(
             loss=loss if loss is not None else None,
-            logits=(rm_logits if self.n_ahead > 1 else logits) if not self.output_logits_at_the_end else logits,
+            logits=logits,
             past_key_values=outputs.past_key_values,
             hidden_states=outputs.hidden_states,
             attentions=outputs.attentions,
@@ -1385,9 +1385,9 @@ class QuietForSequenceClassification(QuietPreTrainedModel):
             use_cache=use_cache,
             output_attentions=output_attentions,
             output_hidden_states=output_hidden_states,
-            return_dict=return_dict,
+            return_dict=True,  # Set return_dict=True
         )
-        hidden_states = transformer_outputs[0]
+        hidden_states = transformer_outputs.last_hidden_state
         logits = self.score(hidden_states)
 
         if input_ids is not None: