Refactoring: move type_handlers.py to _src/serialization

PiperOrigin-RevId: 693341305
AI-Hypercomputer · Nov 5, 2024 · 0445bc1 · 0445bc1
1 parent ea39d05
commit 0445bc1
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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