diff --git a/src/guidellm/data/loaders.py b/src/guidellm/data/loaders.py
index fbeaf0583..e6393707e 100644
--- a/src/guidellm/data/loaders.py
+++ b/src/guidellm/data/loaders.py
@@ -119,6 +119,8 @@ def generator(
                         # This should be fixed at some point.
                         row = preprocessor(row)  # type: ignore[assignment]
                     yield row  # type: ignore[misc]
+                except StopIteration:
+                    raise  # Stop iteration when any dataset is exhausted
                 except Exception as err:  # noqa: BLE001 # Exception logged
                     logger.error(f"Skipping data row due to error: {err}")
                     gen_count -= 1