formatting

facebookresearch · Jan 4, 2023 · 135972e · 135972e
1 parent 2f7e693
commit 135972e
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/metaseq/model_parallel/modules/sequence_parallel_transformer_layer.py b/metaseq/model_parallel/modules/sequence_parallel_transformer_layer.py
@@ -195,7 +195,7 @@ def forward(
                     op=xf_op[0],
                 )
                 .transpose(0, 1)
-                .reshape(seq_len, bsz, num_heads*head_dim)
+                .reshape(seq_len, bsz, num_heads * head_dim)
             )
             # TODO: Reshape q/k/v back to original?
         else:
@@ -413,7 +413,11 @@ def backward(ctx, grad_output):
                 op=xf_op[0],
             )
             out = attn
-            attn = attn.transpose(0, 1).reshape(seq_len, bsz, num_heads*head_dim).contiguous()
+            attn = (
+                attn.transpose(0, 1)
+                .reshape(seq_len, bsz, num_heads * head_dim)
+                .contiguous()
+            )
         else:
             attn, attn_probs = SequeuceParallelTransformerBlock.forward_mha(
                 q, k, v, bsz, seq_len, head_dim, embed_dim_per_partition, dtype