Handle _io_unflatten_object when _thread_local.output_dir is not avai…

…lable Signed-off-by: Hemil Desai <[email protected]>
NVIDIA · Nov 14, 2024 · c506a03 · c506a03
1 parent 02f0932
commit c506a03
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/nemo/lightning/io/mixin.py b/nemo/lightning/io/mixin.py
@@ -608,7 +608,9 @@ def _io_flatten_object(instance):
 
 
 def _io_unflatten_object(values, metadata):
-    assert hasattr(_thread_local, "output_dir")
+    if not hasattr(_thread_local, "output_dir"):
+        return fdl.Config.__unflatten__(values, metadata)
+
     output_dir = _thread_local.output_dir
 
     if len(values) == 1: