Update modeling_quiet.py

modeling_quiet.py

@@ -955,6 +955,11 @@ class QuietModel(QuietPreTrainedModel):
 
             thought_ids.append(thought_id.squeeze(0))
             thought_embeddings.append(thought_embedding.squeeze(0))
+            seq_length = hidden_states.size(1)
+            thought_embeddings = [
+                torch.nn.functional.pad(emb, (0, 0, 0, seq_length - emb.size(0)), mode='constant', value=0)[:seq_length]
+                for emb in thought_embeddings
+    ]
 
         return thought_ids, thought_embeddings