def _accumulate_next_gradients(self):
    'Accumulates the gradients from the next element in dataset.'
    iterator = iter(self.train_dataset)

    @tf.function
    def _accumulate_next():
        (per_replica_features, per_replica_labels) = next(iterator)
        return self._accumulate_gradients(per_replica_features, per_replica_labels)
    while True:
        try:
            (yield _accumulate_next())
        except tf.errors.OutOfRangeError:
            break
Accumulates the gradients from the next element in dataset.
def _accumulate_next_gradients(self):
    
    iterator = iter(self.train_dataset)

    @tf.function
    def _accumulate_next():
        (per_replica_features, per_replica_labels) = next(iterator)
        return self._accumulate_gradients(per_replica_features, per_replica_labels)
    while True:
        try:
            (yield _accumulate_next())
        except tf.errors.OutOfRangeError:
            break