diff --git a/tests/pytorch/r2.1/hf-diffusers.libsonnet b/tests/pytorch/r2.1/hf-diffusers.libsonnet
index 5141ff6ca..8bc7667a9 100644
--- a/tests/pytorch/r2.1/hf-diffusers.libsonnet
+++ b/tests/pytorch/r2.1/hf-diffusers.libsonnet
@@ -24,6 +24,7 @@ local tpus = import 'templates/tpus.libsonnet';
       '--lr_warmup_steps=0',
       '--output_dir=/tmp/sd-pokemon-model',
       '--checkpoints_total_limit=3',
+      '--checkpointing_steps=6000',
     ],
   },