Correct shape (#1170)

huggingface · Jan 4, 2024 · dc53b8c · dc53b8c
1 parent 20428c4
commit dc53b8c
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/trl/trainer/ppo_trainer.py b/trl/trainer/ppo_trainer.py
@@ -1045,15 +1045,15 @@ def train_minibatch(
 
         Args:
             logprobs (`torch.FloatTensor`):
-                Log probabilities of the model, shape [batch_size, response_length]
+                Log probabilities of the model, shape [mini_batch_size, response_length]
             values (`torch.FloatTensor`):
-                Values of the value head, shape [batch_size, response_length]
+                Values of the value head, shape [mini_batch_size, response_length]
             query (`torch.LongTensor`):
-                Encoded queries, shape [batch_size, query_length]
+                Encoded queries, shape [mini_batch_size, query_length]
             response (`torch.LongTensor`):
-                Encoded responses, shape [batch_size, response_length]
+                Encoded responses, shape [mini_batch_size, response_length]
             model_input (`torch.LongTensor`):
-                Concatenated queries and responses, shape [batch_size, query_length+response_length]
+                Concatenated queries and responses, shape [mini_batch_size, query_length+response_length]
 
         Returns:
             train_stats (dict[str, `torch.Tensor`]):