work around for pytorch stft backward bug.

[boltzmann-Li]

All I did was moving the STFT computation outside the autocasting context, and added cast to float to make sure it's using float. Will also file an issue in pytorch repo.