Internal change.

PiperOrigin-RevId: 557561606
google-research · Aug 16, 2023 · dd1b2a8 · dd1b2a8
1 parent ec630b0
commit dd1b2a8
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/t5x/checkpoints.py b/t5x/checkpoints.py
@@ -2393,7 +2393,7 @@ def restore(
     restore_kwargs = {
         _STATE_KEY: {
             'restore_args': restore_args,
-            'transform_fn': transform_fn,
+            'legacy_transform_fn': transform_fn,
         },
     }
     restored = self._manager.restore(