Update modeling_quiet.py

modeling_quiet.py

@@ -1431,30 +1431,28 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
                 # Apply the language model head to get the final logits
                 logits = self.lm_head(mixed_hidden_states)
 
-                if not return_dict:
-                    return logits
-
-                return BaseModelOutputWithPast(
-                    logits=logits,
-                    past_key_values=new_key_values,
-                    hidden_states=outputs_after.hidden_states if output_hidden_states else None,
-                    attentions=outputs_after.attentions if output_attentions else None,
-                )
-            else:
-                # Handle the case where next_token_id is not defined (e.g., continuation_length <= 0)
-                # This part of the code needs to be adapted based on how you want to handle this scenario.
-                # As a placeholder, returning the logits from the last state of the original input.
-                logits = self.lm_head(hidden_states_before)
-
-                if not return_dict:
+                if return_dict:
+                    return BaseModelOutputWithPast(
+                        logits=logits,
+                        past_key_values=new_key_values,
+                        hidden_states=outputs_after.hidden_states if output_hidden_states else None,
+                        attentions=outputs_after.attentions if output_attentions else None,
+                    )
+                else:
                     return logits
-
-                return BaseModelOutputWithPast(
-                    logits=logits,
-                    past_key_values=past_key_values,
-                    hidden_states=outputs_before.hidden_states if output_hidden_states else None,
-                    attentions=outputs_before.attentions if output_attentions else None,
-                )
+                else:
+                    # Handle the case where next_token_id is not defined (e.g., continuation_length <= 0)
+                    logits = self.lm_head(hidden_states_before)
+
+                    if return_dict:
+                        return BaseModelOutputWithPast(
+                            logits=logits,
+                            past_key_values=past_key_values,
+                            hidden_states=outputs_before.hidden_states if output_hidden_states else None,
+                            attentions=outputs_before.attentions if output_attentions else None,
+                        )
+                    else:
+                        return logits
 
     @torch.no_grad()
     def generate(