From ff876fec403a995c97a7f25c213e68fb58f2c4b4 Mon Sep 17 00:00:00 2001 From: "wangnan39@huawei.com" Date: Fri, 19 Jun 2020 15:22:07 +0800 Subject: [PATCH] fix loop size error when exec train after model init --- mindspore/train/dataset_helper.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/mindspore/train/dataset_helper.py b/mindspore/train/dataset_helper.py index cf09e3a067..14797e568b 100644 --- a/mindspore/train/dataset_helper.py +++ b/mindspore/train/dataset_helper.py @@ -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