fix

Signed-off-by: Ivan Butygin <[email protected]>
iree-org · Dec 13, 2024 · ebc543f · ebc543f
1 parent 8096c00
commit ebc543f
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/iree/turbine/kernel/wave/scheduling/modulo_scheduling.py b/iree/turbine/kernel/wave/scheduling/modulo_scheduling.py
@@ -154,8 +154,10 @@ def schedule_graph(self) -> tuple[dict[fx.Node, int], bool]:
                 break
         else:
             raise Exception("Failed to schedule the graph.")
-        pool.close()
-        pool.join()
+
+        if pool is not None:
+            pool.close()
+            pool.join()
 
         self._initiation_interval = T
         return self.schedule, success