This repository is set up to showcase the fix proposed in: tensorflow/tensorflow#31578
git clone https://github.com/DEKHTIARJonathan/KerasLossScaleOptimizer_Demo.git && cd KerasLossScaleOptimizer_Demo
Download Link: https://drive.google.com/drive/folders/1ryQGhurGivPV-N-16m_jPIpnedfm4QR0?usp=sharing
docker build -t keras_lossscaleoptimizer_demo .
docker run --runtime=nvidia -it --rm \
--network=host \
-e NVIDIA_VISIBLE_DEVICES=0 \
-v $(pwd):/project \
keras_lossscaleoptimizer_demo bash
./update_keras.sh && python keras_amp_example.py