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