import os
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel


def save_checkpoint(path, state_dict, tmp_suffix=".tmp"):
    tmp = path + tmp_suffix
    torch.save({k: v.half().cpu() for k, v in state_dict.items()}, tmp)
    os.replace(tmp, path)


def load_checkpoint(path, map_location="cpu"):
    return torch.load(path, map_location=map_location)


def iter_module_state_dict(model):
    raw = model.module if isinstance(model, DistributedDataParallel) else model
    raw = getattr(raw, '_orig_mod', raw)
    return raw.state_dict()


def unwrap(model):
    raw = model.module if isinstance(model, DistributedDataParallel) else model
    return getattr(raw, '_orig_mod', raw)