Lyken17 / pytorch-OpCounter

Count the MACs / FLOPs of your PyTorch model.
MIT License
4.9k stars 528 forks source link

adding kwarg inputs for model #219

Open aninda-leonardo opened 9 months ago

aninda-leonardo commented 9 months ago

We often need to supply keyword arguments for model forward pass.

It's probably also worth returning the model outputs along the flops and param counts, but I did not want to change the output signature in this one.