Method step

tests/unit/runtime/sparse_tensor/test_sparse_grads.py:36–42

(self, closure=None)

    @torch.no_grad()
    def step(self, closure=None):
        loss_1 = self.adam.step(closure)
        loss_2 = self.adam_sparse.step(closure)

        if loss_1 is not None and loss_2 is not None:
            return loss_1 + loss_2
        return loss_1 or loss_2


class TestSparseAdam(DistributedTest):