    def clip_grad_norm(self, max_norm, aggregate_norm_fn=None):
        """Clips gradient norm."""
        self.scaler.unscale_(self.optimizer)
        grad_norm = self.fp32_optimizer.clip_grad_norm(max_norm, aggregate_norm_fn)
        if not torch.isfinite(grad_norm).all():
            new_loss_scale = self.next_loss_scale
            if new_loss_scale <= self.min_loss_scale:
                raise FloatingPointError(
                    (
                        "AMP: Minimum loss scale reached ({}). Your loss is probably exploding. "
                        "Try restarting training or use fp32. {}"
                    ).format(self.min_loss_scale, new_loss_scale)
                )
            else:
                logger.info(
                    "AMP: overflow detected, setting scale to " f"to {new_loss_scale}"
                )
        return grad_norm
Clips gradient norm.
fairseq/optim/amp_optimizer.py