From c52ea9a253c809c5de3bf1f7efeb35061ff5c615 Mon Sep 17 00:00:00 2001 From: andrewwhitecdw Date: Wed, 12 Aug 2026 15:03:08 -0500 Subject: [PATCH] fix: simplify qkv_format checks using a membership tuple Squashed to single commit for review. Original PR: https://github.com/andrewwhitecdw/TransformerEngine/pull/13 --- tests/pytorch/attention/run_attention_with_cp.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/pytorch/attention/run_attention_with_cp.py b/tests/pytorch/attention/run_attention_with_cp.py index 7c6cdefd15..a83e140483 100644 --- a/tests/pytorch/attention/run_attention_with_cp.py +++ b/tests/pytorch/attention/run_attention_with_cp.py @@ -451,7 +451,7 @@ def run_dpa_with_cp( for x in [q_orig, k_orig, v_orig, dout_orig] + ([] if bias is None else [bias]) ] bias_ = rest[0] if len(rest) else None - if qkv_format == "bshd" or qkv_format == "sbhd": + if qkv_format in ("bshd", "sbhd"): seq_dim = qkv_format.index("s") q_, k_, v_, dout_ = [ x.view( @@ -584,7 +584,7 @@ def run_dpa_with_cp( out, dq, dk, dv, dbias, out_, dq_, dk_, dv_, dbias_ = tensors ############ compare results between CP and no-CP ############ - if qkv_format == "bshd" or qkv_format == "sbhd": + if qkv_format in ("bshd", "sbhd"): if is_training: dq, dk, dv, out = [ x.view(