Internal change.

PiperOrigin-RevId: 555241590
google-research · Aug 9, 2023 · 932d987 · 932d987
1 parent 97d3bbf
commit 932d987
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/t5x/checkpoints.py b/t5x/checkpoints.py
@@ -2391,7 +2391,7 @@ def restore(
     restore_kwargs = {
         _STATE_KEY: {
             'restore_args': restore_args,
-            'transform_fn': transform_fn,
+            'legacy_transform_fn': transform_fn,
         },
     }
     restored = self._manager.restore(