[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
NVIDIA · Nov 5, 2024 · f50f6f0 · f50f6f0
1 parent 0f0dd3d
commit f50f6f0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py
@@ -1709,7 +1709,7 @@ def jax_cond_wrap():
                         q_seq_offsets,
                         k_seq_offsets,
                         seed,
-                        config=helper.get_step_config(config.attn_mask_type)
+                        config=helper.get_step_config(config.attn_mask_type),
                     )
 
                 softmax_aux = helper.correct_softmax_aux(softmax_aux, softmax_aux_per_step)