class_docstr: Run training in mixed precision
class_docstr_tok: ["Run","training","in","mixed","precision"]
code: def after_step(self):
      self.model.zero_grad()  # Zero the gradients of the model manually (optimizer disconnected)
      to_model_params(self.model_pgs, self.master_pgs, self.flat_master)
  
imports: from __future__ import annotations
  from ..basics import *
  from .progress import *
  from torch.cuda.amp import GradScaler, autocast
  from torch.cuda.amp.grad_scaler import OptState
  from ..fp16_utils import convert_network, model_grads_to_master_grads, master_params_to_model_params
  from torch.nn.utils import parameters_to_vector
label_desc: Use image node_id 7 for calling the NonNativeMixedPrecision obj's underlying member method code with example usage: obj.after_step() without return types