@tf.function
def _accumulate_next():
    (per_replica_features, per_replica_labels) = next(iterator)
    return self._accumulate_gradients(per_replica_features, per_replica_labels)
@tf.function
def _accumulate_next():
    (per_replica_features, per_replica_labels) = next(iterator)
    return self._accumulate_gradients(per_replica_features, per_replica_labels)