Method zero_grad

codegeex/megatron/optimizer/optimizer.py:464–467

Copied from torch.optim.optimizer

(self, set_to_none=True)

        self._scale = torch.cuda.FloatTensor([1.0])

    def zero_grad(self, set_to_none=True):
        """Copied from torch.optim.optimizer"""
        for group in self.optimizer.param_groups:
            _zero_grad_group_helper(group["params"], set_to_none)

    def get_loss_scale(self):
        """FP32 optimizer does not do any scaling."""