def METHOD_NAME(function, *args):
    """Checkpointining without re-entrant autograd
    Args:
        function: describes what to run in the forward pass of the model or
            part of the model. It should also know how to handle the inputs
            passed as the tuple. For example, in LSTM, if user passes
            ``(activation, hidden)``, :attr:`function` should correctly use the
            first input as ``activation`` and the second input as ``hidden``
        *args: Arguments to pass in to the given ``function``.
    """
    storage: List[Union[flow.Tensor, None]] = []
    counter = 0
    def pack(x):
        nonlocal counter
        counter += 1
        return counter - 1
    # TODO(jianhao): support restoring rng state once we have flow.random.fork_rng
    def unpack(x):
        if len(storage) == 0:
            def inner_pack(inner):
                storage.append(inner)
                return None
            def inner_unpack(packed):
                raise RuntimeError(
                    "You are calling backwards on a tensor that is never exposed. Please open an issue."
                )
            with flow.enable_grad():
                with flow.autograd.graph.saved_tensors_hooks(inner_pack, inner_unpack):
                    _unused = function(*args)
        return storage[x]
    with flow.autograd.graph.saved_tensors_hooks(pack, unpack):
        output = function(*args)
    return output