message: Strategy should be initialized in the same mode as what the test is run in
diff: @@ -144,7 +144,7 @@ def _augment_with_special_arguments(test_method):
  """A wrapped test method that treats some arguments in a special way."""
  mode = kwargs.pop("mode", "graph")
  - distribution = kwargs.pop("distribution", None)
  + distribution = kwargs.get("distribution", None)
  required_tpu = kwargs.pop("required_tpu", False)
  required_gpus = kwargs.pop("required_gpus", None)
  @@ -153,7 +153,6 @@ def _augment_with_special_arguments(test_method):
  "Do not use `required_gpus` and `distribution` together.")
  assert required_tpu is False, (
  "Do not use `required_tpu` and `distribution` together.")
  - kwargs["distribution"] = distribution.strategy
  required_gpus = distribution.required_gpus
  required_tpu = distribution.required_tpu
  @@ -189,9 +188,13 @@ def _augment_with_special_arguments(test_method):
  if mode == "eager":
  with ops.Graph().as_default(), context.eager_mode():
  + if distribution:
  + kwargs_to_pass["distribution"] = distribution.strategy
  test_method(**kwargs_to_pass)
  elif mode == "graph":
  with ops.Graph().as_default(), context.graph_mode():
  + if distribution:
  + kwargs_to_pass["distribution"] = distribution.strategy
  test_method(**kwargs_to_pass)
  else:
  raise ValueError(
  