Fix select_k_features + units

pysr/sr.py

@@ -1521,9 +1521,6 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
             if X_units is not None:
                 X_units = [X_units[i] for i in self.selection_mask_]
                 self.X_units_ = copy.deepcopy(X_units)
-            if y_units is not None:
-                y_units = [y_units[i] for i in self.selection_mask_]
-                self.y_units_ = copy.deepcopy(y_units)
 
             # Re-perform data validation and feature name updating
             X, y = self._validate_data(X=X, y=y, reset=True, multi_output=True)

pysr/test/test.py

@@ -925,6 +925,7 @@ class TestDimensionalConstraints(unittest.TestCase):
             unary_operators=["cos"],
             **self.default_test_kwargs,
             early_stop_condition=1e-8,
+            select_k_features=3,
         )
         model.fit(self.X, y, X_units=["m", "m", "m", "m", "m"], y_units="m")