github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-25 11:44:03 +00:00
Code Issues Packages Projects Releases Wiki Activity

[engine] fixed bug in gradient accumulation dataloader to keep the last step (#1030)

Browse Source
This commit is contained in:
Frank Lee
2022-05-26 14:28:23 +08:00
committed by GitHub
parent 32291dd73f
commit e4685832f8
1 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
colossalai/engine/gradient_accumulation/_gradient_accumulation.py
View File

@@ -145,6 +145,7 @@ class GradAccumDataloader:
def __next__(self) -> Union[Tensor, Tuple[Tensor]]:
if self._cur_step < self.steps_per_epoch:
self._cur_step += 1
data = next(self._dataiter)
if self._cur_step == self.steps_per_epoch and self.consume_remain_data:
# this is to handle non standard pytorch dataloader
@@ -154,7 +155,7 @@ class GradAccumDataloader:
_ = next(self._dataiter)
except StopIteration:
break
return next(self._dataiter)
return data
else:
raise StopIteration