if self.cache is None:
            self.cache = [_zero(g.shape) for g in grads]