Update modeling_quiet.py

modeling_quiet.py

@@ -1682,18 +1682,16 @@ class QuietForCausalLM(QuietPreTrainedModel):
                     assert self.no_residual
                     residual_logits = self.lm_head(hidden_states)
                     talk_hidden_states = hidden_states
+                if 'hidden_states_lm' not in locals():
+                    hidden_states_lm = hidden_states
+                rm_hidden_states = hidden_states
+                if ahead_idx > self.n_ahead - 1:
+                    cur_base_hidden = torch.cat([
+                        base_hidden_states[..., ahead_idx - self.n_ahead + 1:, :],
+                        base_hidden_states[..., :ahead_idx - self.n_ahead + 1, :]
+                    ], dim=-2)
                 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([
-                            base_hidden_states[..., ahead_idx - self.n_ahead + 1:, :],
-                            base_hidden_states[..., :ahead_idx - self.n_ahead + 1, :]
-                        ], dim=-2)
-                    else:
-                        cur_base_hidden = base_hidden_states
+                    cur_base_hidden = base_hidden_states
 
                     if self.use_concat_talk_head:
                         # concatenate the hidden states with the original hidden states