[BC-breaking] Change dtype of softmax to support TorchScript and MyPy (#68336)
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/68336
Test Plan: Imported from OSS
Reviewed By: VitalyFedyunin
Differential Revision: D32470965
Pulled By: cpuhrsch
fbshipit-source-id: 254b62db155321e6a139bda9600722c948f946d3
