diff --git a/examples/imagenet/ddp_analyze.py b/examples/imagenet/ddp_analyze.py
index 2499190..e9a0c8c 100644
--- a/examples/imagenet/ddp_analyze.py
+++ b/examples/imagenet/ddp_analyze.py
@@ -53,13 +53,13 @@ def parse_args():
     parser.add_argument(
         "--query_batch_size",
         type=int,
-        default=128,
+        default=64,
         help="Batch size for computing query gradient.",
     )
     parser.add_argument(
         "--train_batch_size",
         type=int,
-        default=128,
+        default=256,
         help="Batch size for computing training gradient.",
     )