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