Merge pull request #847 from deepmodeling/zjgemi

fix: multiprocessing in debug mode
deepmodeling · Aug 7, 2024 · da9ca61 · da9ca61
2 parents 09fb5d2 + fe6fdd2
commit da9ca61
Showing 1 changed file with 2 additions and 1 deletion.
diff --git 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: