Solved by baldassarre.fe in post #9 
              
              
                I like this solution! 
To add my 50 cents, I would use numel() instad of np.prod() and compress the expression in one line: 
def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)