From b9742de3a3f13a8690b7ecf719c37d49fa486094 Mon Sep 17 00:00:00 2001
From: ms_yan <6576637+ms_yan@user.noreply.gitee.com>
Date: Fri, 17 Apr 2020 15:24:10 +0800
Subject: [PATCH] repeair GeneratorDataset to_device problem

---
 mindspore/dataset/engine/datasets.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/mindspore/dataset/engine/datasets.py b/mindspore/dataset/engine/datasets.py
index 4cff9c57b9..77b7dc54c7 100644
--- a/mindspore/dataset/engine/datasets.py
+++ b/mindspore/dataset/engine/datasets.py
@@ -607,9 +607,9 @@ class Dataset:
 
         def get_distribution(output_dataset):
             dev_id = 0
-            if isinstance(output_dataset, (StorageDataset, GeneratorDataset, MindDataset)):
+            if isinstance(output_dataset, (StorageDataset, MindDataset)):
                 return output_dataset.distribution, dev_id
-            if isinstance(output_dataset, (Cifar10Dataset, Cifar100Dataset, ImageFolderDatasetV2,
+            if isinstance(output_dataset, (Cifar10Dataset, Cifar100Dataset, GeneratorDataset, ImageFolderDatasetV2,
                                            ManifestDataset, MnistDataset, VOCDataset, CelebADataset)):
                 sampler = output_dataset.sampler
                 if isinstance(sampler, samplers.DistributedSampler):