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

    return model.state_dict()