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)
175 Likes