diff --git a/aimodel/src/subcommands/train.py b/aimodel/src/subcommands/train.py
index 3d0b214..0969451 100644
--- a/aimodel/src/subcommands/train.py
+++ b/aimodel/src/subcommands/train.py
@@ -55,7 +55,7 @@ def run(args):
 		feature_dim_in=args.feature_dim,
 		
 		metadata = read_metadata(args.input),
-		shape_water=[ args.water_size, args.water_size ] # The DESIRED output shape. the actual data will be cropped to match this.
+		shape_water_out=[ args.water_size, args.water_size ] # The DESIRED output shape. the actual data will be cropped to match this.
 	)
 	
 	ai.train(dataset_train, dataset_validate)