def set_state_dict(model, state_dict):
    """Load state dictionary whether or not distributed training was used"""
    if hasattr(model, "module"):
        return model.module.load_state_dict(state_dict)

    return model.load_state_dict(state_dict)