EstimatorSpec: Define default scaffold in case of None provided. This will reduce the clutter in the usage of EstimatorSpec. Forexample, We can check scaffold.saver instead of if scaffold and scaffold.saver...

Change: 149482422

1 files changed, 8 insertions, 0 deletions

diff --git a/tensorflow/python/estimator/model_fn_test.py b/tensorflow/python/estimator/model_fn_test.py
index cf6608ae4c..29ca79ff51 100644
--- a/tensorflow/python/estimator/model_fn_test.py
+++ b/tensorflow/python/estimator/model_fn_test.py
@@ -182,6 +182,14 @@ class EstimatorSpecTrainTest(test.TestCase):
             train_op=control_flow_ops.no_op(),
             scaffold=_InvalidScaffold())
 
+  def testReturnDefaultScaffold(self):
+    with ops.Graph().as_default(), self.test_session():
+      estimator_spec = model_fn.EstimatorSpec(
+          mode=model_fn.ModeKeys.FIT,
+          loss=constant_op.constant(1.),
+          train_op=control_flow_ops.no_op())
+      self.assertIsNotNone(estimator_spec.scaffold)
+
 
 class EstimatorSpecEvalTest(test.TestCase):
   """Tests EstimatorSpec in eval mode."""