fix pool bug

AI-Hypercomputer · Nov 27, 2024 · 4795c97 · 4795c97
1 parent 03f5ef6
commit 4795c97
Showing 1 changed file with 5 additions and 12 deletions.
diff --git a/pathwaysutils/persistence/pathways_orbax_handler.py b/pathwaysutils/persistence/pathways_orbax_handler.py
@@ -165,23 +165,16 @@ async def deserialize(
             )
 
             grouped_arrays_and_futures = None
-            with Pool() as p:
-                grouped_arrays_and_futures = p.map(
-                    f(
-                        location=location,
-                        name=name,
-                        dtype=dtype,
-                        shape=shape,
-                        shardings=sharding,
-                    )
-                    for location, name, dtype, shape, sharding in zip(
+            from functools import partial
+            args_list = [[location, name, dtype, shape, sharding] for location, name, dtype, shape, sharding in zip(
                         locations,
                         names,
                         grouped_dtypes,
                         grouped_global_shapes,
                         grouped_shardings,
-                    )
-                )
+                    )]
+            with Pool() as p:
+                grouped_arrays_and_futures = p.apply(f, args_list)
             # grouped_arrays_and_futures = [
             #     f(
             #         location=location,