message: PSv2: Apply strategy.run() change to parameter_server_training_test.
diff: @@ -146,6 +146,8 @@ class KPLTest(test.TestCase):
  @def_function.function
  def worker_fn(iterator):
  +
  + def train_step(iterator):
  batch_data, labels = next(iterator)
  with backprop.GradientTape() as tape:
  pred = model(batch_data, training=True)
  @@ -159,6 +161,8 @@ class KPLTest(test.TestCase):
  actual_pred = math_ops.cast(math_ops.greater(pred, 0.5), dtypes.int64)
  accuracy.update_state(labels, actual_pred)
  + self.client._strategy.run(train_step, args=(iterator,))
  +
  distributed_iterator = iter(distributed_dataset)
  for _ in range(10):
  self.client.schedule(worker_fn, args=(distributed_iterator,))
  