fix size of return matrix in solve

src/ode_solver/method.rs

@@ -9,7 +9,7 @@ use crate::{
     AdjointContext, AdjointEquations, AugmentedOdeEquations, Checkpointing, DefaultDenseMatrix,
     DenseMatrix, Matrix, NewtonNonlinearSolver, NonLinearOp, OdeEquations, OdeEquationsAdjoint,
     OdeEquationsSens, OdeSolverProblem, OdeSolverState, Op, SensEquations, StateRef, StateRefMut,
-    VectorViewMut,
+    Vector, VectorViewMut,
 };
 
 use super::checkpointing::HermiteInterpolator;
@@ -122,7 +122,6 @@ where
         self.set_problem(state, problem)?;
         let mut ret_t = Vec::new();
         let mut ret_y = Vec::new();
-        let nstates = problem.eqn.rhs().nstates();
         let mut write_out = |t: Eqn::T, y: &Eqn::V, g: &Eqn::V| {
             ret_t.push(t);
             match problem.eqn.out() {
@@ -159,8 +158,8 @@ where
             self.state().unwrap().g,
         );
         let ntimes = ret_t.len();
-        let mut ret_y_matrix =
-            <<Eqn::V as DefaultDenseMatrix>::M as Matrix>::zeros(nstates, ntimes);
+        let nrows = ret_y[0].len();
+        let mut ret_y_matrix = <<Eqn::V as DefaultDenseMatrix>::M as Matrix>::zeros(nrows, ntimes);
         for (i, y) in ret_y.iter().enumerate() {
             ret_y_matrix.column_mut(i).copy_from(y);
         }