def save_grads(model_dict, t): 
    """
    A closure to save the gradient wrt to input of a nn module.
    Arguments:
    - model_dict: defaultdict(list)
    - t: dictionary key (usually training iteration)
    """
    def hook(module, grad_input, grad_output):
        model_dict[t].append(grad_input[0].data)

    return hook