From 516c4fe7b2740e2e5a68e6050aeb2fb2737345a0 Mon Sep 17 00:00:00 2001 From: andrewwhitecdw Date: Thu, 13 Aug 2026 17:36:01 -0500 Subject: [PATCH] fix: guard softmax_offset.grad.zero_() by is_training Signed-off-by: andrewwhitecdw --- tests/pytorch/attention/run_attention_with_cp.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/pytorch/attention/run_attention_with_cp.py b/tests/pytorch/attention/run_attention_with_cp.py index 7c6cdefd15..1f578ecb59 100644 --- a/tests/pytorch/attention/run_attention_with_cp.py +++ b/tests/pytorch/attention/run_attention_with_cp.py @@ -517,7 +517,7 @@ def run_dpa_with_cp( torch.cuda.Stream(), cp_comm_type, ) - if config.softmax_type != "vanilla": + if is_training and config.softmax_type != "vanilla": core_attn.softmax_offset.grad.zero_() if dtype == "fp8": core_attn.fp8_initialized = False