提交 fb763b70 编写于 作者: H HydrogenSulfate

fix(dali): reset dataloader manually when using DALI dataloader during epoch training

上级 1c68b63a
......@@ -30,6 +30,9 @@ def train_epoch(engine, epoch_id, print_batch_step):
try:
batch = next(engine.train_dataloader_iter)
except Exception:
# NOTE: reset DALI dataloader manually
if engine.use_dali:
engine.train_dataloader.reset()
engine.train_dataloader_iter = iter(engine.train_dataloader)
batch = next(engine.train_dataloader_iter)
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册