From 6343c14e52fbc0b9a4d34a5897248f3e1dd5b8c4 Mon Sep 17 00:00:00 2001 From: Colin Gaffney Date: Wed, 3 Jul 2024 13:21:20 -0700 Subject: [PATCH] Stop writing msgpack file for new checkpoints and update empty nodes handling so that it no longer depends on this file. PiperOrigin-RevId: 649179323 --- vmoe/initialization/initialization.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vmoe/initialization/initialization.py b/vmoe/initialization/initialization.py index cb3a621..b778b9c 100644 --- 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(