def METHOD_NAME(self, strategy, optimizer_fn):
    coordinator = tf.distribute.experimental.coordinator.ClusterCoordinator(
        strategy
    )
    with strategy.scope():
        model = self._get_model()
        optimizer = optimizer_fn()
        def per_worker_dataset_fn():
            return strategy.distribute_datasets_from_function(
                self._get_dataset_fn()
            )
        ds = coordinator.create_per_worker_dataset(per_worker_dataset_fn)
        @tf.function
        def train_step(iterator):
            def replica_fn(data):
                features, labels = data
                with tf.GradientTape() as tape:
                    output = model(tf.expand_dims(features, axis=1))
                    loss = keras.losses.MeanSquaredError(
                        reduction=losses_utils.ReductionV2.NONE
                    )(labels, output)
                grads = tape.gradient(loss, model.trainable_variables)
                optimizer.apply_gradients(
                    zip(grads, model.trainable_variables)
                )
            strategy.run(replica_fn, args=(next(iterator),))
        for _ in range(3):
            coordinator.schedule(train_step, args=(iter(ds),))
            coordinator.join()
        self.assertEqual(self.evaluate(optimizer.iterations), 3)
        self._verify_accumulators_updated(optimizer)