def METHOD_NAME(model_dir, device):
    checkpoint = torch.load(
        '{}/{}'.format(model_dir, ModelFile.TORCH_MODEL_BIN_FILE),
        map_location=device)
    state_dict = checkpoint['state_dict'].copy()
    for k in checkpoint['state_dict']:
        if k.startswith('avg_model.'):
            v = state_dict.pop(k)
            state_dict[k[4:]] = v
    return state_dict