提交 9a35f409 编写于 作者: D dengkaipeng

stop on return

上级 ef5c4197
...@@ -304,11 +304,9 @@ class Reader(object): ...@@ -304,11 +304,9 @@ class Reader(object):
if self._epoch < 0: if self._epoch < 0:
self.reset() self.reset()
if self.drained(): if self.drained():
self.stop()
raise StopIteration raise StopIteration
batch = self._load_batch() batch = self._load_batch()
if self._drop_last and len(batch) < self._batch_size: if self._drop_last and len(batch) < self._batch_size:
self.stop()
raise StopIteration raise StopIteration
if self._worker_num > -1: if self._worker_num > -1:
return batch return batch
...@@ -419,8 +417,10 @@ def create_reader(cfg, max_iter=0, global_cfg=None, devices_num=1): ...@@ -419,8 +417,10 @@ def create_reader(cfg, max_iter=0, global_cfg=None, devices_num=1):
yield _batch yield _batch
n += 1 n += 1
if max_iter > 0 and n == max_iter: if max_iter > 0 and n == max_iter:
reader.stop()
return return
if max_iter <= 0: if max_iter <= 0:
reader.stop()
return return
reader.reset() reader.reset()
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册