diff --git a/aimodel/src/lib/ai/model_rainfallwater_segmentation.py b/aimodel/src/lib/ai/model_rainfallwater_segmentation.py
index bed0e68..16aedaa 100644
--- a/aimodel/src/lib/ai/model_rainfallwater_segmentation.py
+++ b/aimodel/src/lib/ai/model_rainfallwater_segmentation.py
@@ -26,7 +26,7 @@ def model_rainfallwater_segmentation(metadata, feature_dim_in, shape_water_out,
 	
 	# layer_next = tf.keras.layers.Reshape((1, 1, feature_dim_in), name="cns.stable_begin.reshape")(layer_next)
 	
-	layer_next = do_convnext_inverse(layer_next, arch_name="convnext_i_tiny")
+	layer_next = do_convnext_inverse(layer_next, arch_name=model_arch)
 	
 	# TODO: An attention layer here instead of a dense layer, with a skip connection perhaps?
 	logger.warning("Warning: TODO implement attention from https://ieeexplore.ieee.org/document/9076883")
diff --git a/aimodel/src/subcommands/train.py b/aimodel/src/subcommands/train.py
index b6f1d1c..2fb7171 100644
--- a/aimodel/src/subcommands/train.py
+++ b/aimodel/src/subcommands/train.py
@@ -19,7 +19,7 @@ def parse_args():
 	parser.add_argument("--reads-multiplier", help="Optional. The multiplier for the number of files we should read from at once. Defaults to 1.5, which means read ceil(NUMBER_OF_CORES * 1.5) files at once. Set to a higher number of systems with high read latency to avoid starving the GPU of data.")
 	parser.add_argument("--water-size", help="The width and height of the square of pixels that the model will predict. Smaller values crop the input more [default: 100].", type=int)
 	parser.add_argument("--water-threshold", help="The threshold at which a water cell should be considered water. Water depth values lower than this will be set to 0 (no water). Value unit is metres [default: 0.1].", type=int)
-	parser.add_argument("--arch", help="Next fo the underlying convnext model to use [default: 0.1].")
+	parser.add_argument("--arch", help="Next fo the underlying convnext model to use [default: convnext_i_xtiny].")
 	
 	
 	return parser