await for future

AI-Hypercomputer · Dec 3, 2024 · d2752ee · d2752ee
1 parent d12a41e
commit d2752ee
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/pathwaysutils/persistence/pathways_orbax_handler.py b/pathwaysutils/persistence/pathways_orbax_handler.py
@@ -154,7 +154,7 @@ async def deserialize(
             grouped_dtypes = [dtypes[idx] for idx in idxs]
             grouped_shardings = [shardings[idx] for idx in idxs]
             locations, names = extract_parent_dir_and_name(grouped_infos)
-            grouped_arrays = helper.read_arrays(
+            grouped_arrays, read_future = helper.read_arrays(
                 locations[0],
                 names,
                 grouped_dtypes,
@@ -163,6 +163,7 @@ async def deserialize(
                 global_mesh.devices,
                 timeout=self._read_timeout,
             )
+            read_future.result()
             # f = functools.partial(
             #     helper.read_one_array,
             #     devices=global_mesh.devices,