code: def conjugate_gradient_method(matrix_vector_operator, loss_gradient, nsteps, rdotr_tol=1e-10):
      
      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
docs: Conjugate gradient algorithm