def count_params(net,only_grad = False):
    """
    Count number of parameters with grad or all parameters in model.
    """
    if only_grad:
        return sum(p.numel() for p in net.parameters() if p.requires_grad)
    return sum(p.numel() for p in net.parameters())