diff --git a/aimodel/src/subcommands/train.py b/aimodel/src/subcommands/train.py
index 999c594..dde76cc 100644
--- a/aimodel/src/subcommands/train.py
+++ b/aimodel/src/subcommands/train.py
@@ -65,7 +65,7 @@ def run(args):
 		batch_size=args.batch_size,
 		feature_dim_in=args.feature_dim,
 		
-		model_arch=args.model_arch,
+		model_arch=args.arch,
 		metadata = read_metadata(args.input),
 		shape_water_out=[ args.water_size, args.water_size ], # The DESIRED output shape. the actual data will be cropped to match this.
 	)