torch.cuda.amp.custom_bwd(bwd)
Helper decorator for backward methods of custom autograd functions
   (subclasses of "torch.autograd.Function"). Ensures that "backward"
   executes with the same autocast state as "forward". See the example
   page for more detail.
class torch.cpu.amp.autocast(enabled=True, dtype=torch.bfloat16, cache_enabled=True)
See "torch.autocast". "torch.cpu.amp.autocast(args...)" is
   equivalent to "torch.autocast("cpu", args...)"
Gradient Scaling
If the forward pass for a particular op has "float16" inputs, the
backward pass for that op will produce "float16" gradients. Gradient
values with small magnitudes may not be representable in "float16".
These values will flush to zero ("underflow"), so the update for the
corresponding parameters will be lost.
To prevent underflow, "gradient scaling" multiplies the network's
loss(es) by a scale factor and invokes a backward pass on the scaled
loss(es).  Gradients flowing backward through the network are then