Add support for torch.cuda.amp

[quickgrid]

Apex installation command pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./ failed for me on windows, so I installed it from, https://anaconda.org/conda-forge/nvidia-apex. After using flag --fp16 True training was far slower than fp16 disabled. I tested it on pytorch 1.6 and stylegan2-pytorch 1.2.6.

I came across this issue which mentions torch.cuda.amp fixes many issues from apex.amp and has many more advantages including windows support. Please add support if possible.

[lucidrains]

@quickgrid Last I checked, amp still has some issues in the DDP setting, so I'm holding off until those issues are resolved

Also, when I tried amp, i didn't get any speed or memory improvements on top of using apex

[quickgrid]

Closing issue as amp should mostly speedup newer Nvidia GPU's with compute capability 7.0+. I have also read other posts and articles that pytorch amp will not provide much benefit with mixed precision on Nvidia 10 series GPU's as they lack tensor cores.

This link is for tensorflow mixed precision, but same should apply for pytorch, https://www.tensorflow.org/guide/mixed_precision.