def METHOD_NAME(original_forward, offload_to_cpu, *args, **kwargs):
    # Autograd Functions in PyTorch work best with positional args, since
    # the backward must return gradients (or None) for every input argument.
    # We can flatten keyword arguments to make this easier.
    kwarg_keys, flat_args = pack_kwargs(*args, **kwargs)
    parent_ctx_dict = {"offload": offload_to_cpu}
    output = CheckpointFunction.apply(
        original_forward, parent_ctx_dict, kwarg_keys, *flat_args
    )
    if isinstance(output, torch.Tensor):
        return output
    else:
        packed_non_tensor_outputs = parent_ctx_dict["packed_non_tensor_outputs"]
        if packed_non_tensor_outputs:
            output = unpack_non_tensors(output, packed_non_tensor_outputs)
        return output