Update modeling_quiet.py

modeling_quiet.py

@@ -1662,7 +1662,7 @@ 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)
+            hidden_states_lm = hidden_states
             logits = self.lm_head(hidden_states_lm)
 
             if ahead_idx == 0:
@@ -1682,6 +1682,10 @@ class QuietForCausalLM(QuietPreTrainedModel):
                     assert self.no_residual
                     residual_logits = self.lm_head(hidden_states)
                     talk_hidden_states = hidden_states
+                else:
+                    if 'hidden_states_lm' not in locals():
+                        hidden_states_lm = hidden_states
+                    rm_hidden_states = hidden_states
                 else:
                     if ahead_idx > self.n_ahead - 1:
                         cur_base_hidden = torch.cat([
@@ -1780,6 +1784,7 @@ class QuietForCausalLM(QuietPreTrainedModel):
                         if not self.comparison_mode and not (self.optimize_lm_head_only_at_start and (self.n_ahead + self.n_ahead_talk > 2)) or self.original_mode:
                             loss_list.append(loss)
                         talk_loss_list.append(nonzero_mean(loss).detach())
+
             
             if not attempted or self.comparison_mode:
                 rm_hidden_states = hidden_states