refactor: load eq_param_index of BundleSolvers correctly

NeuroDiffGym · Dec 31, 2023 · b5d457d · b5d457d
1 parent 1b6d8c5
commit b5d457d
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/neurodiffeq/solvers_utils.py b/neurodiffeq/solvers_utils.py
@@ -530,7 +530,9 @@ def load(cls,
                          t_min=t_min,
                          t_max=t_max,
                          theta_min=tuple(load_dict['solver'].r_min[1:]),
-                         theta_max=tuple(load_dict['solver'].r_max[1:]))
+                         theta_max=tuple(load_dict['solver'].r_max[1:]),
+                         eq_param_index=load_dict['solver'].eq_param_index
+            )
 
         if best_nets != None:
             solver.best_nets = best_nets