diff --git a/distributed_shampoo/examples/ddp_cifar10_example.py b/distributed_shampoo/examples/ddp_cifar10_example.py
index 55ce7260..6140ccb4 100644
--- a/distributed_shampoo/examples/ddp_cifar10_example.py
+++ b/distributed_shampoo/examples/ddp_cifar10_example.py
@@ -140,7 +140,7 @@
         raise ValueError(
             "Distributed checkpointing is only supported with DistributedShampoo!"
         )
-    if args.se_distributed_checkpoint and args.checkpoint_dir is None:
+    if args.use_distributed_checkpoint and args.checkpoint_dir is None:
         raise ValueError(
             "Trying to use distributed checkpointing but checkpoint directory is not provided!"
         )