def METHOD_NAME(self):
    """Verify that metric data is complete during evaluate when using
    ParameterServerStrategy
    """
    strategy = tf.distribute.ParameterServerStrategy(
        self.cluster_resolver,
        variable_partitioner=None,
    )
    class BatchCount(metrics_module.Sum):
        def __init__(self, name="batch_count", dtype=tf.int64):
            super().__init__(name=name, dtype=dtype)
        def update_state(self, y_true, y_pred, sample_weight=None):
            return super().update_state(1, sample_weight)
    # Build and compile model within strategy scope.
    with strategy.scope():
        inputs = layers_module.Input((1,))
        outputs = layers_module.Dense(1)(inputs)
        model = training_module.Model(inputs, outputs)
        model.compile(
            loss="mse", metrics=[BatchCount()], steps_per_execution=2
        )
    BATCH_SIZE = 10
    x, y = np.ones((400, 1)), np.ones((400, 1))
    dataset = tf.data.Dataset.from_tensor_slices((x, y))
    batch_count = x.shape[0] // BATCH_SIZE
    # Verify that Model Eval batch counter metric is accurate.
    eval_results = model.evaluate(dataset, steps=batch_count)
    self.assertEqual(eval_results[-1], batch_count)