Automatic Mixed Precision package - torch.amp
"torch.amp" provides convenience methods for mixed precision, where
some operations use the "torch.float32" ("float") datatype and other
operations use lower precision floating point datatype
("lower_precision_fp"): "torch.float16" ("half") or "torch.bfloat16".
Some ops, like linear layers and convolutions, are much faster in
"lower_precision_fp". Other ops, like reductions, often require the
dynamic range of "float32".  Mixed precision tries to match each op to
its appropriate datatype.
Ordinarily, "automatic mixed precision training" with datatype of
"torch.float16" uses "torch.autocast" and "torch.cuda.amp.GradScaler"
together, as shown in the CUDA Automatic Mixed Precision examples and
CUDA Automatic Mixed Precision recipe. However, "torch.autocast" and
"torch.cuda.amp.GradScaler" are modular, and may be used separately if
desired. As shown in the CPU example section of "torch.autocast",