Function tensor_pickling_context

python/oneflow/framework/check_point_v2.py:463–474

(
    path: Path,
    global_rank: Optional[int],
    mp: MAP_LOCATION,
    save_as_external_data: bool,
)

@contextlib.contextmanager
def tensor_pickling_context(
    path: Path,
    global_rank: Optional[int],
    mp: MAP_LOCATION,
    save_as_external_data: bool,
):
    global context_data
    context_data = ContextData(path, global_rank, mp, save_as_external_data)
    try:
        yield
    finally:
        context_data = None


def is_oneflow_pickle_file(path: FILE_LIKE, support_pytorch_format: bool) -> bool: