Update modeling_quiet.py

modeling_quiet.py

@@ -2293,13 +2293,15 @@ class QuietForCausalLM(QuietPreTrainedModel):
         return input_ids
         
     def prepare_thought_embeds(self, hidden_states, temperature=1.0):
+        batch_size, seq_len, hidden_size = hidden_states.shape
+
         if self.use_start_thought_token:
-            start_embed = self.start_embedding[0].unsqueeze(0) * temperature
+            start_embed = self.start_embedding[0].unsqueeze(0).unsqueeze(0).repeat(batch_size, 1, 1) * temperature
         else:
-            start_embed = hidden_states[:, 0, :]
+            start_embed = hidden_states[:, :1, :]
             
         if self.use_end_thought_token:  
-            end_embed = self.end_embedding[0].unsqueeze(0) * temperature
+            end_embed = self.end_embedding[0].unsqueeze(0).unsqueeze(0).repeat(batch_size, 1, 1) * temperature
             thought_embeds = torch.cat([start_embed, hidden_states[:, 1:-1, :], end_embed], dim=1)
         else:
             thought_embeds = torch.cat([start_embed, hidden_states[:, 1:, :]], dim=1)