diff --git a/.github/workflows/pytest.yml b/.github/workflows/pytest.yml
index c360420284d..305759027b5 100644
--- a/.github/workflows/pytest.yml
+++ b/.github/workflows/pytest.yml
@@ -59,7 +59,7 @@ jobs:
         ports:
           - 9000:9000
 
-    timeout-minutes: 90
+    timeout-minutes: 120
     steps:
       - name: Setup ludwigai/ludwig-ray container for local testing with act.
         if: ${{ env.ACT }}
diff --git a/ludwig/automl/defaults/combiner/tabnet_config.yaml b/ludwig/automl/defaults/combiner/tabnet_config.yaml
index 7fa67159356..9966e143135 100644
--- a/ludwig/automl/defaults/combiner/tabnet_config.yaml
+++ b/ludwig/automl/defaults/combiner/tabnet_config.yaml
@@ -16,10 +16,10 @@ hyperopt:
     trainer.learning_rate:
       space: choice
       categories: [0.005, 0.01, 0.02, 0.025]
-    trainer.decay_rate:
+    trainer.learning_rate_scheduler.decay_rate:
       space: choice
       categories: [0.8, 0.9, 0.95]
-    trainer.decay_steps:
+    trainer.learning_rate_scheduler.decay_steps:
       space: choice
       categories: [500, 2000, 8000, 10000, 20000]
     combiner.size:
diff --git a/ludwig/benchmarking/examples/process_config.py b/ludwig/benchmarking/examples/process_config.py
index a137a63391e..8d4db9cd8ff 100644
--- a/ludwig/benchmarking/examples/process_config.py
+++ b/ludwig/benchmarking/examples/process_config.py
@@ -36,9 +36,9 @@ def process_config(ludwig_config: dict, experiment_dict: dict) -> dict:
                     "categories": ["fill_with_const", "fill_with_mean"],
                 },
                 "combiner.type": {"space": "choice", "categories": ["tabnet", "concat"]},
-                "trainer.decay": {"space": "choice", "categories": [True, False]},
+                "trainer.learning_rate_scheduler.decay": {"space": "choice", "categories": [True, False]},
                 "trainer.learning_rate": {"space": "loguniform", "lower": 0.0001, "upper": 0.1},
-                "trainer.decay_rate": {"space": "uniform", "lower": 0.4, "upper": 0.96},
+                "trainer.learning_rate_scheduler.decay_rate": {"space": "uniform", "lower": 0.4, "upper": 0.96},
                 "trainer.batch_size": {"space": "randint", "lower": 32, "upper": 2048},
             },
             "search_alg": {"type": "hyperopt"},
diff --git a/tests/integration_tests/test_automl.py b/tests/integration_tests/test_automl.py
index 89089ffaec8..9c62dbcbc25 100644
--- a/tests/integration_tests/test_automl.py
+++ b/tests/integration_tests/test_automl.py
@@ -27,8 +27,14 @@
 ray = pytest.importorskip("ray")
 
 import dask.dataframe as dd  # noqa
+from ray.tune.experiment.trial import Trial  # noqa
 
-from ludwig.automl import create_auto_config, create_auto_config_with_dataset_profile, train_with_config  # noqa
+from ludwig.automl import (  # noqa
+    auto_train,
+    create_auto_config,
+    create_auto_config_with_dataset_profile,
+    train_with_config,
+)
 from ludwig.hyperopt.execution import RayTuneExecutor  # noqa
 
 pytestmark = pytest.mark.distributed
@@ -299,6 +305,21 @@ def test_train_with_config(time_budget, test_data_tabular_large, ray_cluster_2cp
     _run_train_with_config(time_budget, test_data_tabular_large, tmpdir)
 
 
+@pytest.mark.distributed
+def test_auto_train(test_data_tabular_large, ray_cluster_2cpu, tmpdir):
+    _, ofeatures, dataset_csv = test_data_tabular_large
+    results = auto_train(
+        dataset=dataset_csv,
+        target=ofeatures[0][NAME],
+        time_limit_s=120,
+        user_config={"hyperopt": {"executor": {"num_samples": 2}}},
+    )
+
+    analysis = results.experiment_analysis
+    for trial in analysis.trials:
+        assert trial.status != Trial.ERROR, f"Error in trial {trial}"
+
+
 @pytest.mark.parametrize("fs_protocol,bucket", [private_param(("s3", "ludwig-tests"))], ids=["s3"])
 def test_train_with_config_remote(fs_protocol, bucket, test_data_tabular_large, ray_cluster_2cpu):
     backend = {