Implement a typestr() for PyGrainCheckpointHandler.

PiperOrigin-RevId: 705160710
google · Dec 11, 2024 · 567ade4 · 567ade4
1 parent 1ed92a1
commit 567ade4
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/grain/_src/python/checkpoint_handlers.py b/grain/_src/python/checkpoint_handlers.py
@@ -89,6 +89,10 @@ def finalize(self, directory: epath.Path):
   def close(self):
     pass
 
+  @classmethod
+  def typestr(cls):
+    return f"{cls.__module__}.{cls.__qualname__}"
+
 
 try:
   # Register the handler to be used with the new checkpointing API if Orbax is