commit 3ef67e3337328209b5d46736a7624366ff2a748a
author A. Unique TensorFlower <gardener@tensorflow.org> Fri Jul 09 11:58:37 2021 -0700
committer TensorFlower Gardener <gardener@tensorflow.org> Fri Jul 09 12:07:04 2021 -0700
tree c118c6e562cc4daf7896977524fe7a31c51ccbc7
parent d5166d805fd6d531245d585d81154e5a375bc0ab

Async wait for checkpoint save and restore.

PiperOrigin-RevId: 383893184
Change-Id: I9334f14d7d4b4c8f0c61d4569d6a351f672ea26f

tensorflow/python/training/tracking/util.py

diff --git a/tensorflow/python/training/tracking/util.py b/tensorflow/python/training/tracking/util.py
index 490afe1..2981eec 100644
--- a/tensorflow/python/training/tracking/util.py
+++ b/tensorflow/python/training/tracking/util.py
@@ -2183,8 +2183,6 @@
         model_checkpoint_path=file_path,
         all_model_checkpoint_paths=[file_path],
         save_relative_paths=True)
-    if not graph_building:
-      context.async_wait()  # Ensure save operations have completed.
     return file_path
 
   def read(self, save_path, options=None):
@@ -2344,8 +2342,6 @@
 
     try:
       status = self.read(save_path, options=options)
-      if context.executing_eagerly():
-        context.async_wait()  # Ensure restore operations have completed.
     except errors_impl.NotFoundError as e:
       raise errors_impl.NotFoundError(
           None, None,