def step(self, batch_size, max_norm=None):
self.fp32_trainer.allreduce_grads()<EOL>step_size = batch_size * self._scaler.loss_scale<EOL>if max_norm:<EOL><INDENT>norm, ratio, is_finite = grad_global_norm(self.fp32_trainer._params,<EOL>max_norm * self._scaler.loss_scale)<EOL>step_size = ratio * step_size<EOL>if self._support_nan_check:<EOL><INDENT>self.fp32_trainer.update(step_size)<EOL>overflow = is_finite.asscalar() < <NUM_LIT:1><EOL><DEDENT>else:<EOL><INDENT>overflow = not np.isfinite(norm.asscalar())<EOL>if not overflow:<EOL><INDENT>self.fp32_trainer.update(step_size)<EOL><DEDENT><DEDENT><DEDENT>else:<EOL><INDENT>if self._support_nan_check:<EOL><INDENT>self.fp32_trainer.update(step_size)<EOL>overflow = self._scaler.has_overflow(self.fp32_trainer._params)<EOL><DEDENT>else:<EOL><INDENT>overflow = self._scaler.has_overflow(self.fp32_trainer._params)<EOL>if not overflow:<EOL><INDENT>self.fp32_trainer.update(step_size)<EOL><DEDENT><DEDENT><DEDENT>self._scaler.update_scale(overflow)<EOL>
Makes one step of parameter update. Should be called after
        `fp16_optimizer.backward()`, and outside of `record()` scope.

        Parameters
        ----------
        batch_size : int
            Batch size of data processed. Gradient will be normalized by `1/batch_size`.
            Set this to 1 if you normalized loss manually with `loss = mean(loss)`.
        max_norm : NDArray, optional, default is None
            max value for global 2-norm of gradients.