search

1adrianb / pytorch-estimate-flops

Estimate/count FLOPS for a given neural network using pytorch
https://www.adrianbulat.com
BSD 3-Clause "New" or "Revised" License
304 stars 22 forks source link

Set out_size in parse_node_info to 0 if node can't be parsed #9

Closed senysenyseny16 closed 3 years ago

senysenyseny16 commented 3 years ago

For some types of modules it is not possible to parse out_size, then we can set out_size to 0.

1adrianb commented 3 years ago

Thanks for the pull request. That makes sense, the fx module (pytorch 1.8 and newer) should likely address most of this cases, thought at least for now not all common cases are supported as its still in experimental state.