reverted import module code

Signed-off-by: ParthM-GitHub <[email protected]>
securefederatedai · May 10, 2023 · e36cffb · e36cffb
1 parent 1fc8d3a
commit e36cffb
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/openfl/experimental/utilities/metaflow_utils.py b/openfl/experimental/utilities/metaflow_utils.py
@@ -201,7 +201,8 @@ def __init__(self, flow):
         self._postprocess()
 
     def _create_nodes(self, flow):
-        tree = ast.parse(getsource(flow)).body
+        module = __import__(flow.__module__)
+        tree = ast.parse(getsource(module)).body
         root = [
             n
             for n in tree