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

for more information, see https://pre-commit.ci
NVIDIA · Dec 16, 2024 · dd9159b · dd9159b
1 parent a8ca89c
commit dd9159b
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/transformer_engine/pytorch/attention.py b/transformer_engine/pytorch/attention.py
@@ -1266,7 +1266,10 @@ def get_alibi(
         _alibi_cache["_bottom_right_diagonal"] = bottom_right_diagonal
         bias_dtype = torch.float32 if bias_dtype is None else bias_dtype
         _alibi_cache["_bias_dtype"] = bias_dtype
-        _alibi_cache["_actual_seqlens_q"], _alibi_cache["_actual_seqlens_kv"] = actual_seqlens_q, actual_seqlens_kv
+        _alibi_cache["_actual_seqlens_q"], _alibi_cache["_actual_seqlens_kv"] = (
+            actual_seqlens_q,
+            actual_seqlens_kv,
+        )
         _alibi_cache["_alibi_bias"] = bias.contiguous().to(dtype=bias_dtype, device="cuda")
         _alibi_cache["_alibi_bias_require_update"] = False