Just made maxiter a parameter in train_cg