debug tensor permute for chatglm in reward

sunzeyeah · May 25, 2023 · 5ec3e3a · 5ec3e3a
1 parent fa5457d
commit 5ec3e3a
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/src/models/reward.py b/src/models/reward.py
@@ -49,9 +49,7 @@ def reward(
         if self.model_type == "glm":
             hidden_states = transformer_outputs.mems[-1]
         elif self.model_type == "chatglm":
-            hidden_states = transformer_outputs[0]
-            seq_len, batch_size, hidden_size = hidden_states.shape
-            hidden_states = hidden_states.view(batch_size, seq_len, hidden_size)
+            hidden_states = transformer_outputs[0].permute(1, 0, 2)
         elif self.model_type == "pangu":
             hidden_states = transformer_outputs[0]
             hidden_states = hidden_states.squeeze(1)