DeepSpeed Apex Automatic Mixed Precision

deepspeed_config = {
    'train_batch_size': BATCH_SIZE,
    'gradient_clipping': GRAD_CLIP_NORM,
    'amp': {
        'enabled': True,
        'opt_level': 'O1',
    }
}