From 3cc08198e268c19fcfb00b7550b3eb28c94f5563 Mon Sep 17 00:00:00 2001 From: Mark McMahon Date: Thu, 31 Aug 2023 16:13:00 +0100 Subject: [PATCH] D --- src/mridle/experiment/architecture.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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,