From 3fd4e7c4d7ab379329cf019ffba1c292f8428f7a Mon Sep 17 00:00:00 2001
From: Madhav <Madhavampire@gmail.com>
Date: Tue, 19 Mar 2024 21:58:37 +0530
Subject: [PATCH] Enhanced Error Handling in load_tensors()

---
 checkpoint.py | 16 +++++++++++++++-
 1 file changed, 15 insertions(+), 1 deletion(-)

diff --git a/checkpoint.py b/checkpoint.py
index 1c6e878..4539aa4 100644
--- a/checkpoint.py
+++ b/checkpoint.py
@@ -104,7 +104,21 @@ def load_tensors(shaped_arrays, directory, mesh_config, tensor_indices=None):
         else:
             fs.append(pool.submit(np.zeros, t.shape, dtype=t.dtype))
     wait(fs)
-    return [f.result() for f in fs]
+    
+    # return [f.result() for f in fs]
+    """
+    Improve error reporting in load_tensors by catching exceptions within the futures-
+    and logging detailed information about the failure.
+    """
+    results = []
+    for future in fs:
+        try:
+            result = future.result()
+            results.append(result)
+        except Exception as e:
+            logger.error(f"Failed to load tensor: {e}")
+            raise
+    return results
 
 
 def path_tuple_to_string(path: tuple) -> str: