def _accumulate(iterable, fn=lambda x, y: x + y):
    """Taken from PyTorch"""
    it = iter(iterable)
    try:
        total = next(it)
    except StopIteration:
        return
    yield total
    for element in it:
        total = fn(total, element)
        yield total
