From 02b80276fe98310f876417f506f24e9157fb71b6 Mon Sep 17 00:00:00 2001
From: Noam Shazeer <noam@google.com>
Date: Thu, 11 Feb 2021 22:12:32 -0800
Subject: [PATCH] Allow for disabling the automatic save on shutdown. This is
 bad for mesh-tensorflow, where the variables to be saved haven't gotten
 updated since the previous checkpoint was written.

PiperOrigin-RevId: 357126243
---
 mesh_tensorflow/transformer/utils.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/mesh_tensorflow/transformer/utils.py b/mesh_tensorflow/transformer/utils.py
index 46be5d32..a6edd7dc 100644
--- a/mesh_tensorflow/transformer/utils.py
+++ b/mesh_tensorflow/transformer/utils.py
@@ -1697,7 +1697,8 @@ def get_estimator(model_type, vocabulary, mesh_shape,
       tpu_config=my_tpu_config,
       session_config=session_config,
       save_checkpoints_steps=save_checkpoints_steps,
-      save_checkpoints_secs=None)
+      save_checkpoints_secs=None,
+      save_checkpoint_on_shutdown=False)
 
   transformer_model = build_model(
       model_type=model_type,