Count number trainable parameters in a pytorch model

pytorch_count_params.py

def pytorch_count_params(model):
  "count number trainable parameters in a pytorch model"
  total_params = sum(reduce( lambda a, b: a*b, x.size()) for x in model.parameters())
  return total_params

You can find reduce in from functools import reduce
I assume
or you can use simpler version (without reduce):
total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)