diff --git a/examples/visualizations/train_viz_data_generator.py b/examples/visualizations/train_viz_data_generator.py
index 8841ffae0de6a30743f983c50becbd4b9b4ce29e..4878085d1ca280e0484c71b4c6c69b22e0fc14f3 100644
--- a/examples/visualizations/train_viz_data_generator.py
+++ b/examples/visualizations/train_viz_data_generator.py
@@ -216,7 +216,7 @@ parser.add_argument(
 )
 parser.add_argument(
     "--samples",
-    "-sa",
+    "-sam",
     help="Sets the number of samples (without replacement) to take from the validation set.",
     default=5000,
     type=int,
@@ -237,7 +237,7 @@ animal_id = args.animal_id
 arena_dims = args.arena_dims
 batch_size = args.batch_size
 bayopt_trials = args.bayopt
-checkpoints = args.Zcheckpoint_path
+checkpoints = args.checkpoint_path
 exclude_bodyparts = tuple(args.exclude_bodyparts.split(","))
 gaussian_filter = args.gaussian_filter
 hparams = args.hyperparameters