Fix bart shape comment (#8423)

huggingface · Nov 9, 2020 · a8339b9 · a8339b9
1 parent 46509d1
commit a8339b9
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/src/transformers/modeling_bart.py b/src/transformers/modeling_bart.py
@@ -585,7 +585,7 @@ def forward(
 
         x = F.dropout(x, p=self.dropout, training=self.training)
 
-        # Convert to Bart output format: (seq_len, BS, model_dim) -> (BS, seq_len, model_dim)
+        # Convert to Bart output format: (BS, seq_len, model_dim) ->  (seq_len, BS, model_dim)
         x = x.transpose(0, 1)
         encoder_hidden_states = encoder_hidden_states.transpose(0, 1)
 

diff --git a/src/transformers/modeling_tf_bart.py b/src/transformers/modeling_tf_bart.py
@@ -570,7 +570,7 @@ def call(
             x = self.layernorm_embedding(x + positions)
         x = tf.nn.dropout(x, rate=self.dropout if training else 0)
 
-        # Convert to Bart output format: (seq_len, BS, model_dim) -> (BS, seq_len, model_dim)
+        # Convert to Bart output format: (BS, seq_len, model_dim) ->  (seq_len, BS, model_dim)
         x = tf.transpose(x, perm=(1, 0, 2))
         assert len(shape_list(encoder_hidden_states)) == 3, "encoder_hidden_states must be a 3D tensor"
         encoder_hidden_states = tf.transpose(encoder_hidden_states, perm=(1, 0, 2))