diff --git a/aimodel/src/lib/ai/model_rainfallwater_segmentation.py b/aimodel/src/lib/ai/model_rainfallwater_segmentation.py
index 95fb646..f46897c 100644
--- a/aimodel/src/lib/ai/model_rainfallwater_segmentation.py
+++ b/aimodel/src/lib/ai/model_rainfallwater_segmentation.py
@@ -18,7 +18,7 @@ def model_rainfallwater_segmentation(metadata, feature_dim_in, shape_water_out,
 	layer_next = tf.keras.layers.ReLU(name="cns.stage_begin.relu1")(layer_next)
 	layer_next = tf.keras.layers.LayerNormalization(name="cns.stage_begin.norm", epsilon=1e-6)(layer_next)
 	
-	layer_next = tf.keras.layers.Reshape((4, 4, feature_dim_in/(4*4)), name="cns.stable_begin.reshape")(layer_next)
+	layer_next = tf.keras.layers.Reshape((4, 4, math.floor(feature_dim_in/(4*4))), name="cns.stable_begin.reshape")(layer_next)
 	layer_next = tf.keras.layers.Dense(name="cns.stage.begin.dense2", units=feature_dim_in)(layer_next)
 	layer_next = tf.keras.layers.ReLU(name="cns.stage_begin.relu2")(layer_next)
 	layer_next = tf.keras.layers.LayerNormalization(name="cns.stage_begin.norm", epsilon=1e-6)(layer_next)