diff --git a/psiflow/__init__.py b/psiflow/__init__.py
index 5ad963d..d662163 100644
--- a/psiflow/__init__.py
+++ b/psiflow/__init__.py
@@ -1,8 +1,16 @@
-import parsl
+# def wait():
+#    print('asdlfkjaslkdf')
+#    import time
+#    time.sleep(3)
+#    #import parsl
+#    #parsl.wait_for_current_tasks()
+#
+# import atexit
+# atexit.register(wait)
 
 from .execution import ExecutionContextLoader, load_from_yaml  # noqa: F401
 
 load = ExecutionContextLoader.load
 parse_config = ExecutionContextLoader.parse_config
 context = ExecutionContextLoader.context
-wait = parsl.wait_for_current_tasks
+wait = ExecutionContextLoader.wait
diff --git a/psiflow/execution.py b/psiflow/execution.py
index 525f0d7..e1a83e5 100644
--- a/psiflow/execution.py
+++ b/psiflow/execution.py
@@ -1,6 +1,5 @@
 from __future__ import annotations  # necessary for type-guarding class methods
 
-import atexit
 import copy
 import logging
 import math
@@ -447,7 +446,6 @@ def load(
         )
         path_context = path / "context_dir"
         cls._context = ExecutionContext(config, definitions, path_context)
-        atexit.register(parsl.wait_for_current_tasks)
         return cls._context
 
     @classmethod
@@ -456,6 +454,10 @@ def context(cls):
             raise RuntimeError("No ExecutionContext is currently loaded")
         return cls._context
 
+    @classmethod
+    def wait(cls):
+        parsl.wait_for_current_tasks()
+
 
 def load_from_yaml(path: Union[str, Path]) -> ExecutionContext:
     assert ExecutionContextLoader._context is None  # no previously loaded context