From 75bc01123ea658ee1165a195f49a915697f8eba7 Mon Sep 17 00:00:00 2001 From: Russell Power Date: Wed, 9 May 2018 10:23:15 -0700 Subject: [PATCH] Fix bug in handling of SAVERS collection for shutdown hook. PiperOrigin-RevId: 195989954 --- tensorflow/contrib/tpu/python/tpu/session_support.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/tensorflow/contrib/tpu/python/tpu/session_support.py b/tensorflow/contrib/tpu/python/tpu/session_support.py index faf677a..3e91e2d 100644 --- a/tensorflow/contrib/tpu/python/tpu/session_support.py +++ b/tensorflow/contrib/tpu/python/tpu/session_support.py @@ -292,14 +292,21 @@ class GracefulShutdownHook(session_run_hook.SessionRunHook): if self._saver: return self._saver - savers = ops.get_collection(ops.GraphKeys.SAVERS)[0] + savers = ops.get_collection(ops.GraphKeys.SAVERS) if not savers: return None if not isinstance(savers, list): return savers - assert len(savers) == 1, 'Only one saver supported.' + if len(savers) > 1: + logging.error( + 'Multiple savers in the SAVERS collection. On-demand checkpointing ' + 'will be disabled. Pass an explicit `saver` to the constructor to ' + 'override this behavior.' + ) + return None + return savers[0] def after_run(self, run_context, run_values): -- 2.7.4