ML Katas

Mixed Precision Training with autocast

medium (<30 mins) pytorch training mixed precision gpu
this year by E

Modify a training loop to use torch.cuda.amp.autocast:

  • Wrap forward + loss in autocast.
  • Use GradScaler for backward.

Compare training speed vs. full precision.