import torch
import torch.nn as nn

def _set_module(model, submodule_key, module):
    tokens = submodule_key.split('.')
    sub_tokens = tokens[:-1]
    cur_mod = model
    for s in sub_tokens:
        cur_mod = getattr(cur_mod, s)
    setattr(cur_mod, tokens[-1], module)

def init_yolo(M):
    for m in M.modules():
        if isinstance(m, nn.BatchNorm2d):
            m.eps = 1e-5
            m.momentum = 0.03

def load_checkpoint(model, ckpt_path):
    ckpt = torch.load(ckpt_path, map_location="cpu")['model']
    if 'module' in list(ckpt.keys())[0]:
        new_ckpt = {}
        for key, val in ckpt.items():
            new_ckpt[key.replace('module.', '')] = val
        ckpt = new_ckpt

    try:
        if hasattr(model, 'module'):
            model.module.load_state_dict(ckpt, strict=True)
        else:
            model.load_state_dict(ckpt, strict=True)
    except RuntimeError as e:
        print(e)

    return model
