Use jit tracing functionality for more accurate capture of all ops

[rwightman]

Current impl based on forward hook only captures operations in the model that have been placed in Modules. If functional ops are used, they are missed.

I believe a solution would involve utilizing the functionality of the JIT tracing capability, by basing the analysis on the ScriptModule or graph output from the trace function.

[ppwwyyxx]

btw, this was implemented in https://github.com/facebookresearch/fvcore/blob/master/fvcore/nn/flop_count.py