fix bug in training model by amp (#4874)
fix bug in training model by amp (#4874)

* fix bug in amp

* fix bug in training by amp

* support use gradient clipping when amp is enabled

Co-authored-by: Vasilis Vryniotis <d748cd3f763058dcf20e5489b9bc2373b78db68c@users.noreply.github.com>
Co-authored-by: Prabhat Roy <de68f51b8671e71a0285845858f23ed328c37e06@fb.com>
pytorch/vision,pytorch/vision,pytorch/vision,pytorch/vision,pytorch/vision,pytorch/vision