Skip to content

Commit 8f1e001

Browse files
authored
Unmask StopIteration in DataLoader (#468)
## Summary <!-- Include a short paragraph of the changes introduced in this PR. If this PR requires additional context or rationale, explain why the changes are necessary. --> Fixes bug where dataloader will not recognize the end of a dataset due to the StopIteration Exception being caught and masked. --- - [x] "I certify that all code in this PR is my own, except as noted below." ## Use of AI - [ ] Includes AI-assisted code completion - [ ] Includes code generated by an AI application - [ ] Includes AI-generated tests (NOTE: AI written tests should have a docstring that includes `## WRITTEN BY AI ##`)
2 parents 1741526 + 513861b commit 8f1e001

File tree

1 file changed

+2
-0
lines changed

1 file changed

+2
-0
lines changed

src/guidellm/data/loaders.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -119,6 +119,8 @@ def generator(
119119
# This should be fixed at some point.
120120
row = preprocessor(row) # type: ignore[assignment]
121121
yield row # type: ignore[misc]
122+
except StopIteration:
123+
raise # Stop iteration when any dataset is exhausted
122124
except Exception as err: # noqa: BLE001 # Exception logged
123125
logger.error(f"Skipping data row due to error: {err}")
124126
gen_count -= 1

0 commit comments

Comments
 (0)