From f67f62d37ba1a7eb4215f373535c85f86c119dbc Mon Sep 17 00:00:00 2001 From: Daniel Ng Date: Tue, 5 Nov 2024 07:24:28 -0800 Subject: [PATCH] Refactoring: move type_handlers.py to _src/serialization PiperOrigin-RevId: 693341305 --- pathwaysutils/persistence/pathways_orbax_handler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pathwaysutils/persistence/pathways_orbax_handler.py b/pathwaysutils/persistence/pathways_orbax_handler.py index 43b20db..0f5cb08 100644 --- a/pathwaysutils/persistence/pathways_orbax_handler.py +++ b/pathwaysutils/persistence/pathways_orbax_handler.py @@ -22,7 +22,7 @@ from absl import logging import jax from orbax.checkpoint import future -from orbax.checkpoint import type_handlers +from orbax.checkpoint._src.serialization import type_handlers from pathwaysutils.persistence import helper ParamInfo = type_handlers.ParamInfo