def get_model_weight_grad(model):
    """Get model weight grad in a dict."""
    param_grad_dict = {name: param.grad for name, param in model.named_parameters() if param.requires_grad}
    return param_grad_dict