Remove another for-loop in SoftMax (#60873)
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/60873
Test Plan: Sandcastle
Reviewed By: ngimel
Differential Revision: D29406429
fbshipit-source-id: 3b5710ed9e5d1d14379f64670638ab119d0d78e3
