Browse Source

fix loop size error when exec train after model init

tags/v0.5.0-beta
wangnan39@huawei.com 5 years ago
parent
commit
ff876fec40
1 changed files with 5 additions and 5 deletions
  1. +5
    -5
      mindspore/train/dataset_helper.py

+ 5
- 5
mindspore/train/dataset_helper.py View File

@@ -83,12 +83,12 @@ class DatasetHelper:
class _DatasetIter: class _DatasetIter:
"""Base iter for dataset help""" """Base iter for dataset help"""
def __init__(self, dataset): 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, '__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.__TRANSFER_DATASET__ = _exec_datagraph(dataset, self.loop_size)
dataset.__ME_INITED__ = dataset.__TRANSFER_DATASET__.queue_name dataset.__ME_INITED__ = dataset.__TRANSFER_DATASET__.queue_name




Loading…
Cancel
Save