diff --git a/.github/workflows/pytest.yml b/.github/workflows/pytest.yml
index eb4f9594cd6..238702fbc3d 100644
--- a/.github/workflows/pytest.yml
+++ b/.github/workflows/pytest.yml
@@ -28,15 +28,15 @@ jobs:
           - python-version: "3.8"
             pytorch-version: 2.0.0
             torchscript-version: 1.10.2
-            ray-version: 2.2.0
+            ray-version: 2.8.1
           - python-version: "3.9"
             pytorch-version: 2.1.1
             torchscript-version: 1.10.2
-            ray-version: 2.3.1
+            ray-version: 2.8.1
           - python-version: "3.10"
             pytorch-version: nightly
             torchscript-version: 1.10.2
-            ray-version: 2.3.1
+            ray-version: 2.8.1
     env:
       PYTORCH: ${{ matrix.pytorch-version }}
       MARKERS: ${{ matrix.test-markers }}
@@ -257,7 +257,7 @@ jobs:
           cat requirements.txt | sed '/^torch[>=<\b]/d' | sed '/^torchtext/d' | sed '/^torchvision/d' | sed '/^torchaudio/d' > requirements-temp && mv requirements-temp requirements.txt
           cat requirements_distributed.txt | sed '/^ray[\[]/d'
           pip install torch==2.0.0 torchtext torchvision torchaudio
-          pip install ray==2.3.0
+          pip install ray==2.8.1
           pip install '.[test]'
           pip list
         shell: bash
@@ -298,7 +298,7 @@ jobs:
           cat requirements.txt | sed '/^torch[>=<\b]/d' | sed '/^torchtext/d' | sed '/^torchvision/d' | sed '/^torchaudio/d' > requirements-temp && mv requirements-temp requirements.txt
           cat requirements_distributed.txt | sed '/^ray[\[]/d'
           pip install torch==2.0.0 torchtext torchvision torchaudio
-          pip install ray==2.3.0
+          pip install ray==2.8.1
           pip install '.[test]'
           pip list
         shell: bash
@@ -374,7 +374,7 @@ jobs:
           pip --version
           python -m pip install -U pip
           pip install torch==2.0.0 torchtext
-          pip install ray==2.3.0
+          pip install ray==2.8.1
           pip install '.'
           pip list
         shell: bash
diff --git a/ludwig/backend/ray.py b/ludwig/backend/ray.py
index bdb7f8bbe8f..1af6db2b0ea 100644
--- a/ludwig/backend/ray.py
+++ b/ludwig/backend/ray.py
@@ -17,6 +17,7 @@
 import contextlib
 import copy
 import logging
+import os
 from functools import partial
 from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union
 
@@ -29,11 +30,11 @@
 from packaging import version
 from ray import ObjectRef
 from ray.air import session
-from ray.air.checkpoint import Checkpoint
-from ray.air.config import DatasetConfig, RunConfig, ScalingConfig
+from ray.air.config import RunConfig, ScalingConfig
 from ray.air.result import Result
+from ray.data import ActorPoolStrategy
+from ray.train._checkpoint import Checkpoint
 from ray.train.base_trainer import TrainingFailedError
-from ray.train.torch import TorchCheckpoint
 from ray.train.trainer import BaseTrainer as RayBaseTrainer
 from ray.tune.tuner import Tuner
 from ray.util.dask import ray_dask_get
@@ -52,6 +53,7 @@
     init_dist_strategy,
     LocalStrategy,
 )
+from ludwig.globals import MODEL_WEIGHTS_FILE_NAME
 from ludwig.models.base import BaseModel
 from ludwig.models.predictor import BasePredictor, get_output_columns, get_predictor_cls
 from ludwig.schema.trainer import ECDTrainerConfig, FineTuneTrainerConfig
@@ -66,10 +68,10 @@
 from ludwig.types import HyperoptConfigDict, ModelConfigDict, TrainerConfigDict, TrainingSetMetadataDict
 from ludwig.utils.batch_size_tuner import BatchSizeEvaluator
 from ludwig.utils.dataframe_utils import is_dask_series_or_df, set_index_name
-from ludwig.utils.fs_utils import get_fs_and_path
+from ludwig.utils.fs_utils import get_fs_and_path, open_file
 from ludwig.utils.misc_utils import get_from_registry
 from ludwig.utils.system_utils import Resources
-from ludwig.utils.torch_utils import initialize_pytorch
+from ludwig.utils.torch_utils import get_torch_device, initialize_pytorch
 from ludwig.utils.types import DataFrame, Series
 
 _ray220 = version.parse(ray.__version__) >= version.parse("2.2.0")
@@ -212,14 +214,26 @@ def train_fn(
             report_tqdm_to_ray=True,
             **executable_kwargs,
         )
-        results = trainer.train(train_shard, val_shard, test_shard, return_state_dict=True, **kwargs)
+        # Results is a tuple object of length 4 that has:
+        #   1. The model state dict
+        #   2. The training statistics
+        #   3. The validation statistics
+        #   4. The test statistics
+        results: tuple = trainer.train(train_shard, val_shard, test_shard, return_state_dict=True, **kwargs)
         torch.cuda.empty_cache()
 
+        # Create a local directory to store checkpoint related data
+        ckpt_dir = os.path.join(kwargs.get("save_path"), "checkpoint")
+        os.makedirs(ckpt_dir, exist_ok=True)
+
+        # Save the state dict to disk and load it back on the main process
+        ckpt_path = os.path.join(ckpt_dir, MODEL_WEIGHTS_FILE_NAME)
+        torch.save(results[0], ckpt_path)
+
         # Passing objects containing Torch tensors as metrics is not supported as it will throw an
         # exception on deserialization, so create a checkpoint and return via session.report() along
-        # with the path of the checkpoint
-        ckpt = Checkpoint.from_dict({"state_dict": results})
-        torch_ckpt = TorchCheckpoint.from_checkpoint(ckpt)
+        # with the path of the checkpoint on disk.
+        ckpt: Checkpoint = Checkpoint.from_directory(ckpt_dir)
 
         # The checkpoint is put in the object store and then retrieved by the Trainable actor to be reported to Tune.
         # It is also persisted on disk by the Trainable (and synced to cloud, if configured to do so)
@@ -229,8 +243,11 @@ def train_fn(
             metrics={
                 "validation_field": trainer.validation_field,
                 "validation_metric": trainer.validation_metric,
+                "train_results": results[1],
+                "val_results": results[2],
+                "test_results": results[3],
             },
-            checkpoint=torch_ckpt,
+            checkpoint=ckpt,
         )
 
     except Exception:
@@ -371,40 +388,6 @@ def __init__(self, trainer_kwargs: Dict[str, Any]) -> None:
             **trainer_kwargs,
         )
 
-    def _get_dataset_configs(
-        self,
-        datasets: Dict[str, Any],
-        stream_window_size: Dict[str, Union[None, float]],
-        data_loader_kwargs: Dict[str, Any],
-    ) -> Dict[str, DatasetConfig]:
-        """Generates DatasetConfigs for each dataset passed into the trainer."""
-        dataset_configs = {}
-        for dataset_name, _ in datasets.items():
-            if _ray230:
-                # DatasetConfig.use_stream_api and DatasetConfig.stream_window_size have been removed as of Ray 2.3.
-                # We need to use DatasetConfig.max_object_store_memory_fraction instead -> default to 20% when windowing
-                # is enabled unless the end user specifies a different fraction.
-                # https://docs.ray.io/en/master/ray-air/check-ingest.html?highlight=max_object_store_memory_fraction#enabling-streaming-ingest # noqa
-                dataset_conf = DatasetConfig(
-                    split=True,
-                    max_object_store_memory_fraction=stream_window_size.get(dataset_name),
-                )
-            else:
-                dataset_conf = DatasetConfig(
-                    split=True,
-                    use_stream_api=True,
-                    stream_window_size=stream_window_size.get(dataset_name),
-                )
-
-            if dataset_name == "train":
-                # Mark train dataset as always required
-                dataset_conf.required = True
-                # Check data loader kwargs to see if shuffle should be enabled for the
-                # train dataset. global_shuffle is False by default for all other datasets.
-                dataset_conf.global_shuffle = data_loader_kwargs.get("shuffle", True)
-            dataset_configs[dataset_name] = dataset_conf
-        return dataset_configs
-
     def run(
         self,
         train_loop_per_worker: Callable,
@@ -417,9 +400,13 @@ def run(
     ) -> Result:
         dataset_config = None
         if dataset is not None:
-            data_loader_kwargs = data_loader_kwargs or {}
-            stream_window_size = stream_window_size or {}
-            dataset_config = self._get_dataset_configs(dataset, stream_window_size, data_loader_kwargs)
+            dataset_config = ray.train.DataConfig(
+                datasets_to_split="all",
+                execution_options=ray.data.ExecutionOptions(
+                    preserve_order=data_loader_kwargs.get("shuffle", True),
+                    verbose_progress=True,
+                ),
+            )
 
         callbacks = callbacks or []
 
@@ -523,17 +510,22 @@ def train(
         self._validation_metric = trainer_results.metrics["validation_metric"]
 
         # Load model from checkpoint
-        ckpt = TorchCheckpoint.from_checkpoint(trainer_results.checkpoint)
-        results = ckpt.to_dict()["state_dict"]
+        ckpt = trainer_results.checkpoint
+
+        with open_file(os.path.join(ckpt.path, MODEL_WEIGHTS_FILE_NAME), "rb") as f:
+            state_dict = torch.load(f, map_location=torch.device(get_torch_device()))
 
         # load state dict back into the model
         # use `strict=False` to account for PEFT training, where the saved state in the checkpoint
         # might only contain the PEFT layers that were modified during training
-        state_dict, *args = results
         self.model.load_state_dict(state_dict, strict=False)
-        results = (self.model, *args)
 
-        return results
+        return (
+            self.model,
+            trainer_results.metrics["train_results"],
+            trainer_results.metrics["val_results"],
+            trainer_results.metrics["test_results"],
+        )
 
     def train_online(self, *args, **kwargs):
         # TODO: When this is implemented we also need to update the
@@ -752,7 +744,7 @@ def batch_predict(
             predictions = dataset.ds.map_batches(
                 batch_predictor,
                 batch_size=self.batch_size,
-                compute="actors",
+                compute=ActorPoolStrategy(),
                 batch_format="pandas",
                 num_cpus=num_cpus,
                 num_gpus=num_gpus,
@@ -1135,7 +1127,7 @@ def batch_transform(self, df: DataFrame, batch_size: int, transform_fn: Callable
             ds = ds.map_batches(
                 transform_fn,
                 batch_size=batch_size,
-                compute="actors",
+                compute=ActorPoolStrategy(),
                 batch_format="pandas",
                 **self._get_transform_kwargs(),
             )
diff --git a/ludwig/data/dataframe/dask.py b/ludwig/data/dataframe/dask.py
index 5f292eeabb6..38aea4ccde6 100644
--- a/ludwig/data/dataframe/dask.py
+++ b/ludwig/data/dataframe/dask.py
@@ -26,7 +26,7 @@
 from dask.diagnostics import ProgressBar
 from packaging import version
 from pyarrow.fs import FSSpecHandler, PyFileSystem
-from ray.data import Dataset, read_parquet
+from ray.data import ActorPoolStrategy, Dataset, read_parquet
 
 from ludwig.api_annotations import DeveloperAPI
 from ludwig.data.dataframe.base import DataFrameEngine
@@ -167,7 +167,7 @@ def map_batches(self, series, map_fn, enable_tensor_extension_casting=True):
 
         with tensor_extension_casting(enable_tensor_extension_casting):
             ds = ray.data.from_dask(series)
-            ds = ds.map_batches(map_fn, batch_format="pandas")
+            ds = ds.map_batches(map_fn, batch_format="pandas", compute=ActorPoolStrategy())
             return ds.to_dask()
 
     def apply_objects(self, df, apply_fn, meta=None):
diff --git a/ludwig/data/dataset/ray.py b/ludwig/data/dataset/ray.py
index 5ad083fa715..e94f404c1cf 100644
--- a/ludwig/data/dataset/ray.py
+++ b/ludwig/data/dataset/ray.py
@@ -19,7 +19,7 @@
 import queue
 import threading
 from functools import lru_cache
-from typing import Dict, Iterable, Iterator, Literal, Optional, Union
+from typing import Dict, Iterable, Literal, Optional, Union
 
 import numpy as np
 import pandas as pd
@@ -29,7 +29,9 @@
 from pyarrow.fs import FSSpecHandler, PyFileSystem
 from pyarrow.lib import ArrowInvalid
 from ray.data import read_parquet
+from ray.data.dataset import Dataset as _Dataset
 from ray.data.dataset_pipeline import DatasetPipeline
+from ray.data.iterator import DataIterator
 
 from ludwig.api_annotations import DeveloperAPI
 from ludwig.backend.base import Backend
@@ -49,6 +51,7 @@
 
 logger = logging.getLogger(__name__)
 
+_ray_240 = version.parse(ray.__version__) >= version.parse("2.4.0")
 _ray_230 = version.parse(ray.__version__) >= version.parse("2.3.0")
 
 
@@ -140,7 +143,8 @@ def initialize_batcher(
         augmentation_pipeline=None,
     ):
         yield RayDatasetBatcher(
-            self.ds.repeat().iter_datasets(),
+            # self.ds is a MaterializedDataset object - the iterator call returns a DataIterator object
+            self.ds.iterator(),
             self.features,
             self.training_set_metadata,
             batch_size,
@@ -234,7 +238,7 @@ def data_format(self):
 class RayDatasetShard(Dataset):
     def __init__(
         self,
-        dataset_shard: DatasetPipeline,
+        dataset_shard: _Dataset,
         features: Dict[str, FeatureConfigDict],
         training_set_metadata: TrainingSetMetadataDict,
     ):
@@ -244,6 +248,10 @@ def __init__(
         self.create_epoch_iter()
 
     def create_epoch_iter(self) -> None:
+        if _ray_240:
+            self.epoch_iter = self.dataset_shard
+            return
+
         if _ray_230:
             # In Ray >= 2.3, session.get_dataset_shard() returns a DatasetIterator object.
             if isinstance(self.dataset_shard, ray.data.DatasetIterator):
@@ -289,7 +297,14 @@ def initialize_batcher(
 
     @lru_cache(1)
     def __len__(self):
-        return next(self.epoch_iter).count()
+        if isinstance(self.epoch_iter, DataIterator):
+            num_rows = 0
+            for block, meta in self.epoch_iter._to_block_iterator()[0]:
+                num_rows += meta.num_rows
+            return num_rows
+        else:
+            # self.epoch_iter is a ray.data.Dataset object
+            return self.epoch_iter.count()
 
     @property
     def size(self):
@@ -306,7 +321,7 @@ def to_scalar_df(self, features: Optional[Iterable[BaseFeature]] = None) -> Data
 class RayDatasetBatcher(Batcher):
     def __init__(
         self,
-        dataset_epoch_iterator: Iterator[DatasetPipeline],
+        dataset_epoch_iterator: _Dataset,
         features: Dict[str, Dict],
         training_set_metadata: TrainingSetMetadataDict,
         batch_size: int,
@@ -364,7 +379,7 @@ def steps_per_epoch(self):
         return math.ceil(self.samples_per_epoch / self.batch_size)
 
     def _fetch_next_epoch(self):
-        pipeline = next(self.dataset_epoch_iterator)
+        pipeline = self.dataset_epoch_iterator
 
         read_parallelism = 1
         if read_parallelism == 1:
@@ -431,14 +446,14 @@ def augment_batch(df: pd.DataFrame) -> pd.DataFrame:
 
         return augment_batch
 
-    def _create_sync_reader(self, pipeline: DatasetPipeline):
+    def _create_sync_reader(self, pipeline: _Dataset):
         def sync_read():
             for batch in pipeline.iter_batches(prefetch_blocks=0, batch_size=self.batch_size, batch_format="pandas"):
                 yield self._prepare_batch(batch)
 
         return sync_read()
 
-    def _create_async_reader(self, pipeline: DatasetPipeline):
+    def _create_async_reader(self, pipeline: _Dataset):
         q = queue.Queue(maxsize=100)
         batch_size = self.batch_size
         augment_batch = self._augment_batch_fn()
@@ -474,7 +489,7 @@ def async_read():
 
         return async_read()
 
-    def _create_async_parallel_reader(self, pipeline: DatasetPipeline, num_threads: int):
+    def _create_async_parallel_reader(self, pipeline: _Dataset, num_threads: int):
         q = queue.Queue(maxsize=100)
 
         batch_size = self.batch_size
diff --git a/ludwig/hyperopt/execution.py b/ludwig/hyperopt/execution.py
index dc479e2a364..115932cc3f4 100644
--- a/ludwig/hyperopt/execution.py
+++ b/ludwig/hyperopt/execution.py
@@ -18,9 +18,10 @@
 
 import ray
 from packaging import version
-from ray import tune
-from ray.air import Checkpoint
+from pyarrow.fs import FileSystem
+from ray import train, tune
 from ray.air.config import CheckpointConfig, FailureConfig, RunConfig
+from ray.train import Checkpoint
 from ray.tune import ExperimentAnalysis, register_trainable, Stopper, TuneConfig
 from ray.tune.execution.placement_groups import PlacementGroupFactory
 from ray.tune.schedulers.resource_changing_scheduler import DistributeResources, ResourceChangingScheduler
@@ -124,8 +125,8 @@ def checkpoint(progress_tracker, save_path):
     def ignore_dot_files(src, files):
         return [f for f in files if f.startswith(".")]
 
-    with tune.checkpoint_dir(step=progress_tracker.tune_checkpoint_num) as checkpoint_dir:
-        checkpoint_model = os.path.join(checkpoint_dir, "model")
+    with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
+        checkpoint_model = os.path.join(temp_checkpoint_dir, "model")
         # Atomic copying of the checkpoints
         if not os.path.isdir(checkpoint_model):
             copy_id = uuid.uuid4()
@@ -385,14 +386,17 @@ def _get_best_model_path(trial_path: str, analysis: ExperimentAnalysis, creds: D
             logger.warning("No best model found")
             yield None
 
-        ckpt_type, ckpt_path = checkpoint.get_internal_representation()
-        if ckpt_type == "uri":
+        ckpt_path = checkpoint.path
+        # The filesystem used by the checkpoint should be a pyarrow filesystem object
+        assert isinstance(checkpoint.filesystem, FileSystem)
+
+        if checkpoint.filesystem.type_name == "local":
+            yield ckpt_path
+        else:
             # Read remote URIs using Ludwig's internal remote file loading APIs, as
             # Ray's do not handle custom credentials at the moment.
             with tempfile.TemporaryDirectory() as tmpdir:
                 yield _download_local_tmpdir(ckpt_path, tmpdir, creds)
-        else:
-            yield ckpt_path
 
     @staticmethod
     def _evaluate_best_model(
@@ -463,8 +467,8 @@ def _run_experiment(
         if "mlflow" in config:
             del config["mlflow"]
 
-        trial_id = tune.get_trial_id()
-        trial_dir = Path(tune.get_trial_dir())
+        trial_id = ray.train.get_context().get_trial_id()
+        trial_dir = Path(ray.train.get_context().get_trial_dir())
 
         modified_config = substitute_parameters(copy.deepcopy(hyperopt_dict["config"]), config)
 
@@ -495,13 +499,16 @@ def report(progress_tracker, split=VALIDATION):
             }
 
             metric_score = tune_executor.get_metric_score(train_stats, split)
-            tune.report(
-                parameters=json.dumps(config, cls=NumpyEncoder),
-                metric_score=metric_score,
-                training_stats=json.dumps(train_stats, cls=NumpyEncoder),
-                eval_stats="{}",
-                trial_id=tune.get_trial_id(),
-                trial_dir=tune.get_trial_dir(),
+            train.report(
+                {
+                    "parameters": json.dumps(config, cls=NumpyEncoder),
+                    "metric_score": metric_score,
+                    "training_stats": json.dumps(train_stats, cls=NumpyEncoder),
+                    "eval_stats": "{}",
+                    "trial_id": ray.train.get_context().get_trial_id(),
+                    "trial_dir": ray.train.get_context().get_trial_dir(),
+                },
+                checkpoint=Checkpoint.from_directory(ray.train.get_context().get_trial_dir()),
             )
 
         class RayTuneReportCallback(Callback):
@@ -657,13 +664,16 @@ def check_queue():
         train_stats, eval_stats = stats.pop()
 
         metric_score = self.get_metric_score(train_stats, hyperopt_dict["eval_split"])
-        tune.report(
-            parameters=json.dumps(config, cls=NumpyEncoder),
-            metric_score=metric_score,
-            training_stats=json.dumps(train_stats, cls=NumpyEncoder),
-            eval_stats=json.dumps(eval_stats, cls=NumpyEncoder),
-            trial_id=tune.get_trial_id(),
-            trial_dir=tune.get_trial_dir(),
+        train.report(
+            {
+                "parameters": json.dumps(config, cls=NumpyEncoder),
+                "metric_score": metric_score,
+                "training_stats": json.dumps(train_stats, cls=NumpyEncoder),
+                "eval_stats": "{}",
+                "trial_id": ray.train.get_context().get_trial_id(),
+                "trial_dir": ray.train.get_context().get_trial_dir(),
+            },
+            checkpoint=Checkpoint.from_directory(ray.train.get_context().get_trial_dir()),
         )
 
     def execute(
@@ -776,12 +786,23 @@ def execute(
             else:
                 search_alg = ConcurrencyLimiter(search_alg, max_concurrent=self.max_concurrent_trials)
 
-        def run_experiment_trial(config, local_hyperopt_dict, checkpoint_dir=None):
+        def run_experiment_trial(config, local_hyperopt_dict):
             # Checkpoint dir exists when trials are temporarily paused and resumed, for e.g.,
             # when using the HB_BOHB scheduler.
+            checkpoint = train.get_checkpoint()
+            if checkpoint:
+                with checkpoint.as_directory() as checkpoint_dir:
+                    return self._run_experiment(
+                        config,
+                        checkpoint_dir,
+                        local_hyperopt_dict,
+                        self.decode_ctx,
+                        _is_ray_backend(backend),
+                    )
+
             return self._run_experiment(
                 config,
-                checkpoint_dir,
+                None,
                 local_hyperopt_dict,
                 self.decode_ctx,
                 _is_ray_backend(backend),
@@ -872,7 +893,7 @@ def _register(name, trainable):
                     ),
                     run_config=RunConfig(
                         name=experiment_name,
-                        local_dir=output_directory,
+                        local_dir=str(output_directory),
                         stop=CallbackStopper(callbacks),
                         callbacks=tune_callbacks,
                         failure_config=FailureConfig(
diff --git a/ludwig/hyperopt/syncer.py b/ludwig/hyperopt/syncer.py
index 940b0fc4830..7b58181c02a 100644
--- a/ludwig/hyperopt/syncer.py
+++ b/ludwig/hyperopt/syncer.py
@@ -1,6 +1,6 @@
 from typing import Any, Callable, Dict, List, Optional, Tuple
 
-from ray.tune.syncer import _BackgroundSyncer
+from ray.train._internal.syncer import _BackgroundSyncer
 
 from ludwig.utils.data_utils import use_credentials
 from ludwig.utils.fs_utils import delete, download, upload
diff --git a/requirements_distributed.txt b/requirements_distributed.txt
index 47f070386ab..730f8074688 100644
--- a/requirements_distributed.txt
+++ b/requirements_distributed.txt
@@ -3,7 +3,7 @@ dask[dataframe]<2023.4.0
 pyarrow
 
 # requirements for ray
-ray[default,data,serve,tune]>=2.2.0,<2.4
+ray[default,data,serve,tune]>=2.8.1
 tensorboardX<2.3
 GPUtil
 tblib
diff --git a/tests/integration_tests/test_ray.py b/tests/integration_tests/test_ray.py
index cd6007cd9ca..1f3c3c2e277 100644
--- a/tests/integration_tests/test_ray.py
+++ b/tests/integration_tests/test_ray.py
@@ -83,7 +83,6 @@
 import ray.exceptions  # noqa: E402
 from ray.air.config import DatasetConfig  # noqa: E402
 from ray.data import Dataset, DatasetPipeline  # noqa: E402
-from ray.train._internal.dataset_spec import DataParallelIngestSpec  # noqa: E402
 
 from ludwig.backend.ray import get_trainer_kwargs, RayBackend  # noqa: E402
 from ludwig.data.dataframe.dask import DaskEngine  # noqa: E402
@@ -1068,6 +1067,8 @@ def create_dataset_pipeline(
             stream_window_size=ds.window_size_bytes,
             global_shuffle=False,
         )
+        from ray.train._internal.dataset_spec import DataParallelIngestSpec  # noqa: E402
+
         spec = DataParallelIngestSpec({"train": dataset_config})
 
         # These two must be called in sequence so that the dataset is tracked internally. No preprocessing is applied.
diff --git a/tests/ludwig/data/test_ray_data.py b/tests/ludwig/data/test_ray_data.py
index a71c8ae910b..2526505a590 100644
--- a/tests/ludwig/data/test_ray_data.py
+++ b/tests/ludwig/data/test_ray_data.py
@@ -26,7 +26,8 @@ def test_async_reader_error():
         "bin1": {},
     }
 
-    with pytest.raises(TypeError, match="'Mock' object is not iterable"):
+    # TODO: See if this is actually the correct fix for this test, not exactly sure what the test is trying to do
+    with pytest.raises(AttributeError, match="'list_iterator' object has no attribute 'iter_batches'"):
         RayDatasetBatcher(
             dataset_epoch_iterator=iter([pipeline]),
             features=features,