Key mismatch when loading NLCDetection weights about hifi_ifdl HOT 2 CLOSED

def restore_weight_helper(model, model_dir, initial_epoch):
    weight_path = f'{model_dir}/{initial_epoch}.pth'
    state_dict = torch.load(weight_path, map_location='cuda:0')['model']

    # Modify the state dict to remove the 'module.' prefix if necessary
    new_state_dict = OrderedDict()
    for k, v in state_dict.items():
        name = k[7:] if k.startswith('module.') else k  # Remove "module."
        new_state_dict[name] = v

    model.load_state_dict(new_state_dict, strict=False)
    print(f'{model_dir} weight-loading succeeds: {weight_path}')

    pytorch_total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"{model_dir}_params: {pytorch_total_params}")
    return model