code: def update(self):
          
  
          # Initialize gradients.
          self.solver.zero_grad()
  
          # Forward and backward
          for _ in range(self.accum_grad):
              # feed data
              self.data_feeder()
  
              # forward
              self.loss.forward(clear_no_need_grad=self.clear_buffer)
  
              # backward with scale
              self.loss.backward(self.scale, clear_buffer=self.clear_buffer)
  
          # AllReduce
          if self.comm and len(self.grads) != 0:
              self.comm.all_reduce(self.grads, division=False, inplace=False)
  
          # Check Inf/NaN in grads
          if self.solver.check_inf_or_nan_grad():
              self.scale /= self.scaling_factor
              self._counter = 0
  
              # Recursively call udpate function until no inf nor nan.
              self._recursive_count += 1
              if self._recursive_count > self._max_recursive_count:
                  self._recursive_count = 0
                  return  # skip
              return self.update()
          self._recursive_count = 0
  
          # Rescale grads
          self.solver.scale_grad(1. / self.scale)
  
          # Do some gradient clipping, etc.
          if self.weight_decay is not None:
              self.solver.weight_decay(self.weight_decay)
  
          # Update
          self.solver.update()
          if self._counter > self.N:
              self.scale *= self.scaling_factor
              self._counter = 0
          self._counter += 1
docs: Monolithic update method.
  
          This method calls the following methods with the dynamic loss scaling.
  
          1. solver.zerograd
          2. feed data
          3. loss.forward
          4. loss.backward
          5. comm.all_reduce (if it is specified)
          6. solver.update