TODO:

  1. Enable AMP (Automatic Mixed Precision) training. Ensure that the CUDA versions of the virtual environment and system are consistent.