未验证 提交 08982c69 编写于 作者: J Jane-Ding 提交者: GitHub

Update program.py (#2665)

Fix training on Windows when one epoch finishes, it quits unexpectedly. Same as fixing the problem when evaluating on Windows.
上级 6445076a
...@@ -197,9 +197,11 @@ def train(config, ...@@ -197,9 +197,11 @@ def train(config,
train_reader_cost = 0.0 train_reader_cost = 0.0
batch_sum = 0 batch_sum = 0
batch_start = time.time() batch_start = time.time()
max_iter = len(train_dataloader) - 1 if platform.system(
) == "Windows" else len(train_dataloader)
for idx, batch in enumerate(train_dataloader): for idx, batch in enumerate(train_dataloader):
train_reader_cost += time.time() - batch_start train_reader_cost += time.time() - batch_start
if idx >= len(train_dataloader): if idx >= max_iter:
break break
lr = optimizer.get_lr() lr = optimizer.get_lr()
images = batch[0] images = batch[0]
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册