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):