[caffe2] Fix shape inference for pruning ops (#53082)
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/53082
Reviewed By: yinghai, khabinov
Differential Revision: D26742532
fbshipit-source-id: 6cdfb293541b601f7916a95e08bf573876c9ca74
