Function save_state

monai/networks/utils.py:627–658

Save the state dict of input source data with PyTorch `save`. It can save `nn.Module`, `state_dict`, a dictionary of `nn.Module` or `state_dict`. And automatically convert the data parallel module to regular module. For example:: save_state(net, path) save_state(net

(src: torch.nn.Module | dict, path: PathLike, **kwargs)

def save_state(src: torch.nn.Module | dict, path: PathLike, **kwargs):
    """
    Save the state dict of input source data with PyTorch `save`.
    It can save `nn.Module`, `state_dict`, a dictionary of `nn.Module` or `state_dict`.
    And automatically convert the data parallel module to regular module.
    For example::

        save_state(net, path)
        save_state(net.state_dict(), path)
        save_state({"net": net, "opt": opt}, path)
        net_dp = torch.nn.DataParallel(net)
        save_state(net_dp, path)

    Refer to: https://pytorch.org/ignite/v0.4.8/generated/ignite.handlers.DiskSaver.html.

    Args:
        src: input data to save, can be `nn.Module`, `state_dict`, a dictionary of `nn.Module` or `state_dict`.
        path: target file path to save the input object.
        kwargs: other args for the `save_obj` except for the `obj` and `path`.
            default `func` is `torch.save()`, details of the args:
            https://pytorch.org/docs/stable/generated/torch.save.html.

    """

    ckpt: dict = {}
    if isinstance(src, dict):
        for k, v in src.items():
            ckpt[k] = get_state_dict(v)
    else:
        ckpt = get_state_dict(src)

    save_obj(obj=ckpt, path=path, **kwargs)


def convert_to_onnx(