Update TfGrainCheckpointHandler.

PiperOrigin-RevId: 686958575
google-research · Oct 17, 2024 · 64b309c · 64b309c
1 parent ce73f91
commit 64b309c
Showing 1 changed file with 1 addition and 9 deletions.
diff --git a/connectomics/jax/checkpoint.py b/connectomics/jax/checkpoint.py
@@ -120,15 +120,7 @@ def restore_checkpoint(
       args=ocp.args.Composite(**restore_args_dict))
 
 
-class TfGrainCheckpointHandler(tfgrain.OrbaxCheckpointHandler):
-
-  def save(self, directory: epath.Path, args: 'TfGrainCheckpointArgs') -> None:
-    return super().save(directory, args.item)
-
-  def restore(
-      self, directory: epath.Path, args: 'TfGrainCheckpointArgs'
-  ) -> tfgrain.TfGrainDatasetIterator:
-    return super().restore(directory, args.item)
+TfGrainCheckpointHandler = tfgrain.OrbaxCheckpointHandler
 
 
 @ocp.args.register_with_handler(  # pytype:disable=wrong-arg-types