optuna for non auto models

georgia-tech-db · Nov 11, 2023 · e03c4d0 · e03c4d0
1 parent 2eb51b2
commit e03c4d0
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/evadb/executor/create_function_executor.py b/evadb/executor/create_function_executor.py
@@ -499,7 +499,8 @@ def get_optuna_config(trial):
                     return model_args_config
 
                 model_args["config"] = get_optuna_config
-                model_args["backend"] = "optuna"
+
+            model_args["backend"] = "optuna"
 
             model_args["h"] = horizon
             model_args["loss"] = MQLoss(level=[conf])
@@ -542,7 +543,7 @@ def get_optuna_config(trial):
                 raise FunctionIODefinitionError(err_msg)
 
             model = StatsForecast(
-                [model_here(season_length=season_length)], freq=new_freq
+                [model_here(season_length=season_length)], freq=new_freq, n_jobs=-1
             )
 
         data["ds"] = pd.to_datetime(data["ds"])

diff --git a/test/integration_tests/long/test_model_forecasting.py b/test/integration_tests/long/test_model_forecasting.py
@@ -110,6 +110,8 @@ def test_forecast(self):
             TYPE Forecasting
             HORIZON 12
             PREDICT 'y'
+            LIBRARY 'neuralforecast'
+            AUTO 'false'
             FREQUENCY 'M';
         """
         execute_query_fetch_all(self.evadb, create_predict_udf)