[pruner] add Conv2d support (#61778)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/61778
Adding Conv2d as supported modules for the pruner. Previously the pruner only supported Linear layers. This addition includes:
- adding a Conv2d activation reconstruction forward hook to match Conv2d weight shapes
- in `prepare`, checking the type of the module and using the corresponding activation forward hook
ghstack-source-id: 134143557
Test Plan:
Added conv2d tests
`buck test mode/dev-nosan //caffe2/test:ao -- TestBasePruner`
https://pxl.cl/1LLf3
Reviewed By: jerryzh168
Differential Revision: D29719045
fbshipit-source-id: 6a9f91b96992c552fff32f0e5a6e22f16eb7077b
