Skip to content

Commit

Permalink
Fixed batch referenced before assignment
Browse files Browse the repository at this point in the history
  • Loading branch information
JAEarly committed Jun 20, 2024
1 parent 523ac9e commit 6d35762
Showing 1 changed file with 10 additions and 0 deletions.
10 changes: 10 additions & 0 deletions composer/trainer/trainer.py
Original file line number Diff line number Diff line change
Expand Up @@ -3644,6 +3644,12 @@ def _iter_dataloader(self, trainer_mode: TrainerMode):
# 0 = not finished, 1 = finished (using integer tensors so we can use dist.all_reduce)
iter_finished = torch.zeros(1, dtype=torch.uint8)
iter_finished = self.state.device.tensor_to_device(iter_finished)

# Initialize batch to avoid "referenced before assignment" warnings
# Unique sentinel value to differentiate uninitialized state and dataloader yielding None
sentinel = object()
batch = sentinel

while True:
try:
# [BEFORE/AFTER]_DATALOADER only runs while training
Expand All @@ -3668,6 +3674,10 @@ def _iter_dataloader(self, trainer_mode: TrainerMode):
if iter_finished.item() == 1:
break

if batch is sentinel:
raise RuntimeError(
"Batch should have been assigned or loop should have been broken. This shouldn't happen!"
)
yield batch

def _use_closures(self) -> bool:
Expand Down

0 comments on commit 6d35762

Please sign in to comment.