From b78efd5b9e7abf8f93cff6a8b320ccce3d67d656 Mon Sep 17 00:00:00 2001
From: panfengfeng <panfengfeng@huawei.com>
Date: Fri, 19 Jun 2020 15:21:45 +0800
Subject: [PATCH] fix gpu invalid input data

---
 .../ccsrc/dataset/engine/datasetops/device_queue_op.cc | 10 +++++++---
 .../ccsrc/dataset/engine/datasetops/device_queue_op.h  |  2 +-
 2 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/mindspore/ccsrc/dataset/engine/datasetops/device_queue_op.cc b/mindspore/ccsrc/dataset/engine/datasetops/device_queue_op.cc
index 18bf659516..84bad9db1a 100644
--- a/mindspore/ccsrc/dataset/engine/datasetops/device_queue_op.cc
+++ b/mindspore/ccsrc/dataset/engine/datasetops/device_queue_op.cc
@@ -232,13 +232,17 @@ Status DeviceQueueOp::RetryPushGPUData(const std::vector<size_t> &data_size, con
 
   while (!GpuBufferMgr::GetInstance().IsClosed() && !TaskManager::FindMe()->Interrupted()) {
     RETURN_IF_NOT_OK(MallocForGPUData(&items, curr_row));
-    auto ret = GpuBufferMgr::GetInstance().Push(handle, items, WAIT_TIME);
+    BlockQueueStatus_T ret = GpuBufferMgr::GetInstance().Push(handle, items, WAIT_TIME);
     if (ret) {
       for (int i = 0; i < items.size(); i++) {
         free(items[i].data_ptr_);
       }
-      MS_LOG(WARNING) << "Retry pushing data...";
-      continue;
+      if (ret == BlockQueueStatus_T::ERROR_INPUT) {
+        return Status(StatusCode::kUnexpectedError, __LINE__, __FILE__, "invalid input Data, please check it.");
+      } else {
+        MS_LOG(WARNING) << "Retry pushing data...";
+        continue;
+      }
     } else {
       break;
     }
diff --git a/mindspore/ccsrc/dataset/engine/datasetops/device_queue_op.h b/mindspore/ccsrc/dataset/engine/datasetops/device_queue_op.h
index 6fec18986e..a854004593 100644
--- a/mindspore/ccsrc/dataset/engine/datasetops/device_queue_op.h
+++ b/mindspore/ccsrc/dataset/engine/datasetops/device_queue_op.h
@@ -25,11 +25,11 @@
 
 #ifdef ENABLE_TDTQUE
 #include "dataset/engine/tdt/tdt_plugin.h"
-
 #endif
 
 #ifdef ENABLE_GPUQUE
 #include "device/gpu/gpu_buffer_mgr.h"
+using mindspore::device::BlockQueueStatus_T;
 using mindspore::device::GpuBufferMgr;
 #endif