From 4896cc21abac49717d4daa1e96b66ab3962d1e56 Mon Sep 17 00:00:00 2001 From: Matt Conley Date: Tue, 16 Apr 2019 11:47:25 -0700 Subject: [PATCH] Fix None gradient test --- tensorflow/python/training/experimental/loss_scale_test.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/tensorflow/python/training/experimental/loss_scale_test.py b/tensorflow/python/training/experimental/loss_scale_test.py index 888abd79ba0..f832022ecb0 100644 --- a/tensorflow/python/training/experimental/loss_scale_test.py +++ b/tensorflow/python/training/experimental/loss_scale_test.py @@ -253,10 +253,9 @@ class DynamicLossScaleTest(test.TestCase, parameterized.TestCase): self._test_helper(inputs, expected_outputs, init_loss_scale) @test_util.run_in_graph_and_eager_modes - def test_update_with_none_gradients(self, strategy_fn): - with strategy_fn().scope(): - loss_scaler = loss_scale_module.DynamicLossScale() - loss_scaler.update([None]) + def test_update_with_none_gradients(self): + loss_scaler = loss_scale_module.DynamicLossScale() + loss_scaler.update([None]) @test_util.run_in_graph_and_eager_modes def test_get(self):