diff --git a/aimodel/src/subcommands/pretrain.py b/aimodel/src/subcommands/pretrain.py
index b05572f..c4e025f 100644
--- a/aimodel/src/subcommands/pretrain.py
+++ b/aimodel/src/subcommands/pretrain.py
@@ -53,7 +53,7 @@ def run(args):
 		feature_dim=args.feature_dim,
 		
 		shape_rainfall=dataset_metadata["rainfallradar"],
-		shape_water=[ math.ceil(value * 0.5) for value in dataset_metadata["waterdepth"] ]
+		shape_water=[ math.ceil(value * 0.5) + 1 for value in dataset_metadata["waterdepth"] ]
 	)
 	
 	ai.train(dataset_train, dataset_validate)