Conjugate gradient algorithm

def conjugate_gradient_method(matrix_vector_operator, loss_gradient, nsteps, rdotr_tol=1e-10):
    """ Conjugate gradient algorithm """
    x = torch.zeros_like(loss_gradient)

    r = loss_gradient.clone()
    p = loss_gradient.clone()

    rdotr = torch.dot(r, r)

    for i in range(nsteps):
        Avp = matrix_vector_operator(p)
        alpha = rdotr / torch.dot(p, Avp)

        x += alpha * p
        r -= alpha * Avp

        new_rdotr = torch.dot(r, r)
        betta = new_rdotr / rdotr
        p = r + betta * p
        rdotr = new_rdotr

        if rdotr < rdotr_tol:
            break

    return x