Docs on `nn`

I think it’s important to document that all ops should be done via nn, else they won’t be tracked (in memory and flops). E.g. x.mean is invisible, as opposed to AvgPool1d (though its flops and memory still aren’t tracked). Would also help to detect if there are non-nn ops, or non-contributing layers like AvgPool1d and warn the user of potential inaccuracies (or document it).

That said, is there a way to ensure correctness for things like x = x + 1? Wrapping torch.add in nn.Module doesn’t do it.

[TylerYep]

Yep, AvgPool is definitely on the todo list. PRs adding support for these are greatly appreciated!

[TylerYep]

Totally agree, a lot of operations are not supported but definitely could in the future. I’ll add this to the documentation.

Note that torchinfo is built around supplementing information of print(model) in Pytorch, so it is unlikely we will list “add” operations and other non-layer operations. But rows that print(model) outputs but are missing from torchinfo should be considered a bug.