Fix issue with dtype being dropped from RestoreArgs.

PiperOrigin-RevId: 615496686
google-research · Mar 13, 2024 · 962b9d5 · 962b9d5
1 parent f6ec080
commit 962b9d5
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/t5x/checkpoints.py b/t5x/checkpoints.py
@@ -2081,7 +2081,7 @@ def _construct_restore_args(
   """Create RestoreArgs for Orbax restoration."""
   if not isinstance(param_info, _OrbaxParamInfo):  # from fallback
     return ocp.RestoreArgs(dtype=dtype)
-  if param_info.name.split('.')[0] != 'target':
+  if param_info.name.split('/')[0] != 'target':
     dtype = None
   if param_info.mesh_axes is None:
     return ocp.RestoreArgs(dtype=dtype)