Avoid saving self for`softmax` and `log_softmax` (#65242)
Summary:
Fixes https://github.com/pytorch/pytorch/issues/64000
- updates double backward formula to compute grad wrt output instead of self
- ~~In some of the error messages, we still refer to the dtype of the input, even though we are now checking the dtype of the output~~
Pull Request resolved: https://github.com/pytorch/pytorch/pull/65242
Reviewed By: malfet
Differential Revision: D31317680
Pulled By: soulitzer
fbshipit-source-id: b3b921e06775cfc12e5a97a9ee8d73aec3aac7c3
