1 files changed, 8 insertions, 1 deletions

diff --git a/pyecsca/misc/utils.py b/pyecsca/misc/utils.py
index aa450e3..970924f 100644
--- a/pyecsca/misc/utils.py
+++ b/pyecsca/misc/utils.py
@@ -47,9 +47,16 @@ class TaskExecutor(ProcessPoolExecutor):
     keys: List[Any]
     futures: List[Future]
 
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.keys = []
+        self.futures = []
+
     def submit_task(self, key: Any, fn, /, *args, **kwargs):
-        self.futures.append(self.submit(fn, *args, **kwargs))
+        future = self.submit(fn, *args, **kwargs)
+        self.futures.append(future)
         self.keys.append(key)
+        return future
 
     @property
     def tasks(self):