RuntimeError and speed loss with opt_level = O1, O2 or O3

Hello,

I discovered your apex tools for integrating mixed precision and FP16 training in pytorch, which is a great idea to develop ! Our servers are mainly equipped with TITAN V cards hence I was really looking forward to trying them out at their fastest. Software versions are pytorch 1.1.0, cuda 9.1.85 and cudnn 7.1.3. As I never tried this before, I used the more straight-forward apex.amp to compare FP32 training with opt_level = O1 or O2 or O3 (the last with keep_batchnorm_fp32=True as my models use batchnorm).

It is a rather large set of codes so I report here my main questions/issues but if needed I can provide more details and can try to give some reproducible issue cases.

#1 the training script I wanted to run with amp uses torch.stft for computing spectral losses
if computing these losses, I get
RuntimeError: arange_out not supported on CPUType for Half
which point to the stft operation. Is that correct that my script should not use spectral operations such as torch.stft to be optimized in mixed precision ? Or is there a fix/workaround for that please ?

#2 I tried to run the comparison only optimizing time domain losses (eg. waveform MSE instead of spectral reconstruction) so that the code runs without error for every opt_level, but then opt_level = O1 or O2 or O3 were all slower than opt_level = O0 (or running my original FP32 training) ... obviously I expected the speed gain to depend on the code, the operations involved, the batch sizes etc. but I did not expect it to be slowed down ...
For this I only used amp.initialize and amp.scale_loss (as recommended in the 1st example of https://nvidia.github.io/apex/amp.html). I train generative models composed mainly of conv1d, batchnorm1d and linear layers. Everything is feed-forward, no softmax or classification. What could I check to understand wether I could hope for speed gain or not in my application case please ?

Good luck developing the mixed precision training, it has a lot of potential if made more integrated in existing tools !

Start by reproducing the reported torch.stft failure with apex.amp at opt_level O1, O2, and O3 using the listed PyTorch, CUDA, and cuDNN versions. Then compare the amp.initialize and amp.scale_loss training path with O0 for the stated convolution, batch-normalization, and linear models; done means identifying the compatibility cause or a confirmed workaround and explaining the speed difference.