diff --git a/aimodel/src/encoderonly_test_rainfall.py b/aimodel/src/encoderonly_test_rainfall.py
index 8e7a657..efb0919 100755
--- a/aimodel/src/encoderonly_test_rainfall.py
+++ b/aimodel/src/encoderonly_test_rainfall.py
@@ -7,6 +7,7 @@ import tensorflow as tf
 
 from lib.dataset.dataset_encoderonly import dataset_encoderonly
 from lib.ai.components.convnext import make_convnext
+from lib.ai.components.CallbackCustomModelCheckpoint import CallbackCustomModelCheckpoint
 from lib.ai.helpers.summarywriter import summarywriter
 
 # ███████ ███    ██ ██    ██