Update modeling_quiet.py

modeling_quiet.py

@@ -1100,10 +1100,10 @@ class QuietForCausalLM(QuietPreTrainedModel):
 
     def _generate_thoughts(self, hidden_states, max_length):
         batch_size = hidden_states.size(0)
-        thought_ids = torch.zeros((batch_size, self.config.num_thoughts, max_length), dtype=torch.long, device=hidden_states.device)
+        thought_ids = torch.zeros((batch_size, self.config.max_thoughts, max_length), dtype=torch.long, device=hidden_states.device)
         thought_embeddings = []
 
-        for i in range(self.config.num_thoughts):
+        for i in range(self.config.max_thoughts):
             thought_input_ids = torch.zeros((batch_size, 1), dtype=torch.long, device=hidden_states.device)
             thought_outputs = self.model.generate(
                 input_ids=thought_input_ids,
@@ -1120,6 +1120,7 @@ class QuietForCausalLM(QuietPreTrainedModel):
         thought_embeddings = torch.stack(thought_embeddings, dim=1)
         return thought_ids, thought_embeddings
 
+
     def calculate_policy_loss(self, thoughts, rewards):
         thought_log_probs = []
         for thought in thoughts: