added checkpointer.wait()

Author: amitsrivastava78

keras/src/callbacks/orbax_checkpoint.py

@@ -468,10 +468,7 @@ def wait_until_finished(self):
         checkpoints if there might be pending save operations.
         """
         # Wait for any async operations to complete
-        while self.checkpointer.is_saving_in_progress():
-            import time
-
-            time.sleep(0.1)
+        self.checkpointer.wait()
 
     def _restore_model_state_from_full_tree(self, state_tree, model=None):
         """Restore model state from full state tree (V1 format)."""