diff --git a/aimodel/src/lib/ai/model_rainfallwater_segmentation.py b/aimodel/src/lib/ai/model_rainfallwater_segmentation.py
index f2b23da..3f0a555 100644
--- a/aimodel/src/lib/ai/model_rainfallwater_segmentation.py
+++ b/aimodel/src/lib/ai/model_rainfallwater_segmentation.py
@@ -14,7 +14,7 @@ def model_rainfallwater_segmentation(metadata, feature_dim_in, shape_water_out,
 	)
 	
 	# BEGIN
-	layer_next = tf.keras.layers.Dense(name="cns.stage.begin.dense")(layer_input)
+	layer_next = tf.keras.layers.Dense(name="cns.stage.begin.dense", units=feature_dim_in)(layer_input)
 	layer_next = tf.keras.layers.LayerNormalisation(name="stage_begin.norm", epsilon=1e-6)(layer_next)
 	layer_next = tf.keras.layers.ReLU(name="stage_begin.relu")(layer_next)