diff --git a/aimodel/src/lib/dataset/dataset_encoderonly.py b/aimodel/src/lib/dataset/dataset_encoderonly.py
index 6ce1783..ea8ae54 100644
--- a/aimodel/src/lib/dataset/dataset_encoderonly.py
+++ b/aimodel/src/lib/dataset/dataset_encoderonly.py
@@ -81,7 +81,7 @@ def parse_item(metadata, water_threshold=0.1, water_bins=2, heightmap=None, rain
 		water = tf.cast(tf.math.greater_equal(water, water_threshold), dtype=tf.int32)
 		
 		rainfall = tf.image.extract_patches(tf.expand_dims(rainfall, axis=0),
-			sizes=[1,windowsize,windowsize],
+			sizes=[1,windowsize,windowsize,1],
 			strides=[1,1,1,1],
 			rates=[1,1,1,1],
 			padding="VALID"