Fix None gradient test

wenkangwei · Apr 16, 2019 · 4896cc2 · 4896cc2
1 parent e73a77d
commit 4896cc2
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/tensorflow/python/training/experimental/loss_scale_test.py b/tensorflow/python/training/experimental/loss_scale_test.py
@@ -253,10 +253,9 @@ def test_random_mix_good_and_bad_gradients(self, strategy_fn):
       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):