Support FP16(amp) in imagenet main.py

[Deeeerek]

I am running test under : https://github.com/pytorch/examples/tree/main/imagenet

Does the main.py support FP16(amp)?like arguments： --amp

Thx

[Steven Liu]

@Deeeerek No. There was a try in #203 I think it should be easier and cleaner to support fp16 nowadays with amp from pytorch: https://pytorch.org/blog/accelerating-training-on-nvidia-gpus-with-pytorch-automatic-mixed-precision/
Please feel free to contribute if you like!

[Deeeerek]

Yes, I did it (amp) with autocast, GradScaler, much thx!