diff --git a/aimodel/src/deeplabv3_plus_test_rainfall.py b/aimodel/src/deeplabv3_plus_test_rainfall.py
index cbe5f83..1af1521 100755
--- a/aimodel/src/deeplabv3_plus_test_rainfall.py
+++ b/aimodel/src/deeplabv3_plus_test_rainfall.py
@@ -153,7 +153,7 @@ history = model.fit(dataset_train,
 			separator="\t"
 		),
 		CallbackCustomModelCheckpoint(
-			model_to_checkpoint=model_predict,
+			model_to_checkpoint=model,
 			filepath=os.path.join(
 				DIR_OUTPUT,
 				"checkpoints"