Method __call__

tests/test_deepspeed_training_worker.py:31–33

(self, *args, **kwargs)

        self.zero_grad_calls = 0

    def __call__(self, *args, **kwargs):
        self.forward_calls += 1
        return self.model(*args, **kwargs)

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