diff --git a/aimodel/src/subcommands/pretrain.py b/aimodel/src/subcommands/pretrain.py
index d1e2737..2d35fe6 100644
--- a/aimodel/src/subcommands/pretrain.py
+++ b/aimodel/src/subcommands/pretrain.py
@@ -1,3 +1,4 @@
+import math
 import sys
 import argparse
 from asyncio.log import logger
@@ -52,7 +53,7 @@ def run(args):
 		feature_dim=args.feature_dim,
 		
 		shape_rainfall=dataset_metadata["rainfallradar"],
-		shape_water=dataset_metadata["waterdepth"]
+		shape_water=[ math.floor(value * 0.75) for value in dataset_metadata["waterdepth"] ]
 	)
 	
 	ai.train(dataset_train, dataset_validate)