diff --git a/src/mridle/experiment/architecture.py b/src/mridle/experiment/architecture.py index d81f8a60..bf324da4 100644 --- a/src/mridle/experiment/architecture.py +++ b/src/mridle/experiment/architecture.py @@ -3,7 +3,7 @@ import skorch from sklearn.base import BaseEstimator from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor -from sklearn.linear_model import LogisticRegression, Lasso +from sklearn.linear_model import LogisticRegression, Lasso, Ridge from sklearn.pipeline import Pipeline from sklearn.compose import ColumnTransformer from sklearn.preprocessing import FunctionTransformer @@ -36,6 +36,7 @@ class ArchitectureInterface(ComponentInterface): 'RandomForestClassifier': RandomForestClassifier, # TODO enable auto-loading from sklearn 'RandomForestRegressor': RandomForestRegressor, # TODO enable auto-loading from sklearn 'LassoRegressor': Lasso, # TODO enable auto-loading from sklearn + 'RidgeRegressor': Ridge, # TODO enable auto-loading from sklearn 'LogisticRegression': LogisticRegression, 'XGBClassifier': xgb.XGBClassifier, 'XGBRegressor': xgb.XGBRegressor,