!3339 [AutoParallel]fix bug for autoparallel gpu

Merge pull request !3339 from lichen/fix_bug_for_autoparallel_gpu
pull/3339/MERGE
mindspore-ci-bot 5 years ago committed by Gitee
commit 26de81a2df

@ -426,7 +426,7 @@ class Model:
# for data sink dataset_helper only iter once, other wise iter epoch_size times.
for inputs in dataset_helper:
if _need_to_full():
if _need_to_full() and context.get_context("device_target") == "GPU":
inputs = _to_full_tensor(inputs, self._device_number, self._global_rank)
list_callback.step_begin(run_context)
outputs = self._train_network(*inputs)

Loading…
Cancel
Save