diff --git a/aimodel/src/lib/dataset/dataset.py b/aimodel/src/lib/dataset/dataset.py
index ff24c90..18b2e02 100644
--- a/aimodel/src/lib/dataset/dataset.py
+++ b/aimodel/src/lib/dataset/dataset.py
@@ -28,8 +28,7 @@ def parse_item(metadata):
 		rainfall = tf.transpose(rainfall, [1, 2, 0])
 		# [width, height] → [width, height, channels]
 		water = tf.expand_dims(water, axis=-1)
-		# BUG: AttributeError: 'dict' object has no attribute 'waterdepth
-		rainfall = tf.image.resize(rainfall, tf.constant(metadata["waterdepth"]))
+		rainfall = tf.image.resize(rainfall, tf.cast(tf.constant(metadata["waterdepth"]) / 2, dtype=tf.int32))
 		
 		# TODO: The shape of the resulting tensor can't be statically determined, so we need to reshape here
 		print("DEBUG:dataset ITEM rainfall:shape", rainfall.shape, "water:shape", water.shape)