Function enable_checkpointing

PATH/core/models/ckpt.py:228–235

Makes :func:`is_checkpointing` return :data:`True` within a context.

()

@contextmanager
def enable_checkpointing() -> Generator[None, None, None]:
    """Makes :func:`is_checkpointing` return :data:`True` within a context."""
    orig = thread_local.is_checkpointing
    thread_local.is_checkpointing = True
    try:
        yield
    finally:
        thread_local.is_checkpointing = orig


@contextmanager