fix loop size error when exec train after model init

This commit is contained in:
wangnan39@huawei.com 2020-06-19 15:22:07 +08:00
parent 3c1b8308cf
commit ff876fec40
1 changed files with 5 additions and 5 deletions

View File

@ -83,12 +83,12 @@ class DatasetHelper:
class _DatasetIter:
"""Base iter for dataset help"""
def __init__(self, dataset):
self.loop_size = 1
if not hasattr(dataset, '__loop_size__'):
self.loop_size = dataset.get_dataset_size()
else:
self.loop_size = dataset.__loop_size__
if not hasattr(dataset, '__ME_INITED__'):
if not hasattr(dataset, '__loop_size__'):
self.loop_size = dataset.get_dataset_size()
else:
self.loop_size = dataset.__loop_size__
dataset.__TRANSFER_DATASET__ = _exec_datagraph(dataset, self.loop_size)
dataset.__ME_INITED__ = dataset.__TRANSFER_DATASET__.queue_name