"automatic mixed precision training/inference" on CPU with datatype of
"torch.bfloat16" only uses "torch.autocast".
For CUDA and CPU, APIs are also provided separately:


"torch.autocast("cuda", args...)" is equivalent to
  "torch.cuda.amp.autocast(args...)".


"torch.autocast("cpu", args...)" is equivalent to
  "torch.cpu.amp.autocast(args...)". For CPU, only lower precision
  floating point datatype of "torch.bfloat16" is supported for now.


Autocasting


Gradient Scaling


Autocast Op Reference


Op Eligibility


CUDA Op-Specific Behavior


CUDA Ops that can autocast to "float16"


CUDA Ops that can autocast to "float32"


CUDA Ops that promote to the widest input type


Prefer "binary_cross_entropy_with_logits" over
  "binary_cross_entropy"




CPU Op-Specific Behavior


CPU Ops that can autocast to "bfloat16"


CPU Ops that can autocast to "float32"


CPU Ops that promote to the widest input type




Autocasting