diff --git a/src/dflow/python/python_op_template.py b/src/dflow/python/python_op_template.py
index ba00b9ee..af0bec21 100644
--- a/src/dflow/python/python_op_template.py
+++ b/src/dflow/python/python_op_template.py
@@ -504,7 +504,8 @@ def render_script(self):
                   " handle_input_parameter\n"
         script += "from dflow.python.utils import handle_output_artifact," \
                   " handle_output_parameter, handle_lineage\n"
-        script += f"from {mod} import {class_name}\n\n"
+        if mod not in ["__main__", "__mp_main__"]:
+            script += f"from {mod} import {class_name}\n\n"
         if hasattr(op_class, "func"):
             script += "op_obj = %s\n" % class_name
         elif op is None: