Nigel
/
mindspore
forked from mindspore-Ecosystem/mindspore
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

fix bug in dataset helper

This commit is contained in:
liyong 2020-12-01 15:42:57 +08:00
parent 8453b0d243
commit 602a8e52a0
1 changed files with 3 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
mindspore/train/dataset_helper.py
View File

@ -106,6 +106,9 @@ def connect_network_with_dataset(network, dataset_helper):
if hasattr(dataset, '__network_manage__') and key in dataset.__network_manage__:
network = dataset.__network_manage__[key]
else:
if _need_to_full():
device_num = _get_device_num()
dataset_shapes = _to_full_shapes(dataset_shapes, device_num)
network = _DataWrapper(network, dataset_types, dataset_shapes, dataset.__transfer_dataset__.queue_name)
dataset.__network_manage__ = dataset.__network_manage__ if hasattr(
dataset, '__network_manage__') else dict()