Fix multi-output index selection

MilesCranmer · Apr 26, 2022 · c4f15ef · c4f15ef
1 parent 20cc2d4
commit c4f15ef
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pysr/sr.py b/pysr/sr.py
@@ -795,7 +795,7 @@ def get_best(self, index=None):
         if index is not None:
             if isinstance(self.equations, list):
                 assert isinstance(index, list)
-                return [self.equations.iloc[i] for i in index]
+                return [eq.iloc[i] for eq, i in zip(self.equations, index)]
             return self.equations.iloc[index]
 
         if self.model_selection == "accuracy":