Stop writing msgpack file for new checkpoints and update empty nodes …

…handling so that it no longer depends on this file. PiperOrigin-RevId: 649179323
google-research · Jul 3, 2024 · 6343c14 · 6343c14
1 parent c0220ef
commit 6343c14
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/vmoe/initialization/initialization.py b/vmoe/initialization/initialization.py
@@ -62,7 +62,7 @@ class PyTreeCheckpointHandlerWithStructure(
     orbax_checkpoint.PyTreeCheckpointHandler):
 
   def structure(self, directory):
-    return self._handler_impl._read_aggregate_file(directory)  # pylint: disable=protected-access
+    return self._read_aggregate_file(directory)  # pylint: disable=protected-access
 
 
 def initialize_from_orbax(