Update modeling_quiet.py

modeling_quiet.py

@@ -1662,6 +1662,9 @@ class QuietForCausalLM(QuietPreTrainedModel):
             prev_rm_logits = rm_logits  # for policy gradient
             prev_rm_tokens = cur_rm_tokens  # for policy gradient
 
+            hidden_states_lm = hidden_states_lm.to(self.lm_head.weight.dtype)
+            logits = self.lm_head(hidden_states_lm)
+
             if ahead_idx == 0:
                 hidden_states_lm = hidden_states
                 logits = self.lm_head(hidden_states_lm)