91b0d220c8e816766fd4565e1d2f5115d3afbefe,gpytorch/utils/cholesky.py,,batch_potrs,#Any#Any#,19

Before Change

    
    
    potrs_list = []
    for i in range(mat.size(0)):
        potrs_list.append(torch.potrs(mat[i], chol[i]).unsqueeze(0))
    return torch.cat(potrs_list, 0)


def tridiag_batch_potrf(trid, upper=False):

After Change

    
    
    potrs_list = []
    potrs_list = [
        torch.potrs(sub_mat, sub_chol)
        for sub_mat, sub_chol in zip(mat.view(-1, *mat.shape[-2:]), chol.view(-1, *chol.shape[-2:]))
    ]
    res = torch.cat(potrs_list, 0)
    return res.view_as(mat)

Instances