diff --git a/metagpt/prompts/ml_engineer.py b/metagpt/prompts/ml_engineer.py
index 33eb9c40c..cca9649b3 100644
--- a/metagpt/prompts/ml_engineer.py
+++ b/metagpt/prompts/ml_engineer.py
@@ -301,6 +301,7 @@ ML_SPECIFIC_PROMPT = {
 ML_MODULE_MAP = {
     "data_preprocess": "metagpt.tools.functions.libs.data_preprocess",
     "feature_engineering": "metagpt.tools.functions.libs.feature_engineering",
+    "udf": "metagpt.tools.functions.libs.udf",
 }
 
 STRUCTURAL_CONTEXT = """