From 9fd76c8693e4d5af0c382ff7b987d0c795f9430c Mon Sep 17 00:00:00 2001
From: Sherlock <baihan.huang@gmail.com>
Date: Mon, 21 Sep 2020 20:21:59 -0700
Subject: [PATCH] Place Shape's output in CPU memory (#5245)

Co-authored-by: Sherlock Huang <bahuang@OrtTrainingDev3.af05slrtruoetgaxwwjv5nsq5e.px.internal.cloudapp.net>
---
 onnxruntime/core/providers/cuda/tensor/shape_op.cc     | 10 ++++++----
 .../core/framework/gradient_graph_builder.h            |  1 -
 2 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/onnxruntime/core/providers/cuda/tensor/shape_op.cc b/onnxruntime/core/providers/cuda/tensor/shape_op.cc
index d50a64d20f..4781c90182 100644
--- a/onnxruntime/core/providers/cuda/tensor/shape_op.cc
+++ b/onnxruntime/core/providers/cuda/tensor/shape_op.cc
@@ -13,8 +13,9 @@ ONNX_OPERATOR_VERSIONED_KERNEL_EX(
     1, 12,
     kCudaExecutionProvider,
     KernelDefBuilder()
-        .OutputMemoryType<OrtMemTypeCPUOutput>(0)
-        .TypeConstraint("T",  DataTypeImpl::AllFixedSizeTensorTypes())
+        // properly force CPU/GPU synch inside the kernel
+        .OutputMemoryType<OrtMemTypeCPUInput>(0)
+        .TypeConstraint("T", DataTypeImpl::AllFixedSizeTensorTypes())
         .TypeConstraint("T1", DataTypeImpl::GetTensorType<int64_t>()),
     Shape);
 
@@ -24,8 +25,9 @@ ONNX_OPERATOR_KERNEL_EX(
     13,
     kCudaExecutionProvider,
     KernelDefBuilder()
-        .OutputMemoryType<OrtMemTypeCPUOutput>(0)
-        .TypeConstraint("T",  DataTypeImpl::AllFixedSizeTensorTypes())
+        // properly force CPU/GPU synch inside the kernel
+        .OutputMemoryType<OrtMemTypeCPUInput>(0)
+        .TypeConstraint("T", DataTypeImpl::AllFixedSizeTensorTypes())
         .TypeConstraint("T1", DataTypeImpl::GetTensorType<int64_t>()),
     Shape);
 
diff --git a/orttraining/orttraining/core/framework/gradient_graph_builder.h b/orttraining/orttraining/core/framework/gradient_graph_builder.h
index ae1f203c4e..eb770e2a53 100644
--- a/orttraining/orttraining/core/framework/gradient_graph_builder.h
+++ b/orttraining/orttraining/core/framework/gradient_graph_builder.h
@@ -60,7 +60,6 @@ static std::unordered_map<std::string, std::unordered_set<size_t>>
         {"Where", {0}},
         {"Range", {0, 1, 2}},
         {"Tile", {1}},
-        {"NonZero", {0}},
         {"BroadcastGradientArgs", {0, 1}}};
 
 class GradientGraphBuilder {