for input, target in data:
          optimizer.zero_grad()
      # Enables autocasting for the forward pass (model + loss)
      with autocast():
          output = model(input)
          loss = loss_fn(output, target)

      # Exits the context manager before backward()
      loss.backward()
      optimizer.step()

See the CUDA Automatic Mixed Precision examples for usage (along
   with gradient scaling) in more complex scenarios (e.g., gradient
   penalty, multiple models/losses, custom autograd functions).
"autocast" can also be used as a decorator, e.g., on the "forward"
   method of your model:
  class AutocastModel(nn.Module):
      ...
      @autocast()
      def forward(self, input):
          ...

Floating-point Tensors produced in an autocast-enabled region may
   be "float16". After returning to an autocast-disabled region, using
   them with floating-point Tensors of different dtypes may cause type