Method zero_grad

tests/test_deepspeed_training_worker.py:35–37

(self, *args, **kwargs)

        return self.model(*args, **kwargs)

    def zero_grad(self, *args, **kwargs):
        self.zero_grad_calls += 1
        self.optimizer.zero_grad(*args, **kwargs)

    def backward(self, loss):
        self.backward_calls += 1