[caffe2] Fix shape inference for Softmax (#53132)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/53132
Input and output should have the same shape for Softmax https://caffe2.ai/docs/operators-catalogue.html#softmax.
Reviewed By: walterddr, yinghai, ChunliF
Differential Revision: D26536592
fbshipit-source-id: 8b50794803aeadcb75d8f370c77f4fef98a1f2ad
