31804a4b46fca8a11e23e22b7d82d10e6cabf582,gpytorch/lazy/cached_cg_lazy_tensor.py,CachedCGLazyTensor,__init__,#CachedCGLazyTensor#Any#Any#Any#,19

Before Change

        // This will make it faster when we reconstruct the LazyTensor inside functions
        with torch.no_grad():
            if solves is None:
                solves = [
                    base_lazy_tensor._solve(eager_rhs, base_lazy_tensor._preconditioner()[0])
                    for eager_rhs in eager_rhss
                ]

        super(CachedCGLazyTensor, self).__init__(
            base_lazy_tensor, eager_rhss=eager_rhss, solves=solves,
        )

After Change

                all_solves = base_lazy_tensor._solve(torch.cat(eager_rhss, -1), base_lazy_tensor._preconditioner()[0])
                solves = []
                for eager_rhs in eager_rhss:
                    solve = all_solves[..., :eager_rhs.size(-1)]
                    all_solves = all_solves[..., eager_rhs.size(-1):]
                    solves.append(solve.detach())  // The detach is necessary here for some reason?

        super(CachedCGLazyTensor, self).__init__(
            base_lazy_tensor, eager_rhss=eager_rhss, solves=solves,

Instances