saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-17 18:40:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Place Shape's output in CPU memory (#5245)

Browse source 
Co-authored-by: Sherlock Huang <bahuang@OrtTrainingDev3.af05slrtruoetgaxwwjv5nsq5e.px.internal.cloudapp.net>
This commit is contained in:
Sherlock 2020-09-21 20:21:59 -07:00 committed by Tianlei Wu
parent 9158679c43
commit 9fd76c8693
2 changed files with 6 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
onnxruntime/core/providers/cuda/tensor/shape_op.cc
View file

@ -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);

1
orttraining/orttraining/core/framework/gradient_graph_builder.h
View file

@ -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 {