Update modeling_quiet.py

modeling_quiet.py

@@ -942,9 +942,10 @@ class QuietModel(QuietPreTrainedModel):
                     inputs_embeds=thought_embedding,
                     attention_mask=None,
                     use_cache=True,
+                    return_dict=True,  # Set return_dict=True
                 )
-                logits = outputs.logits[:, -1, :]
-                next_token_id = torch.argmax(logits, dim=-1)
+                logits = self.lm_head(outputs.last_hidden_state)  # Use outputs.last_hidden_state instead of outputs.logits
+                next_token_id = torch.argmax(logits[:, -1, :], dim=-1)
                 
                 if next_token_id == self.config.end_token_id:
                     break