From c00697cc9e8fed457e4b644062e8c4fcce72f02c Mon Sep 17 00:00:00 2001 From: Kellen Sunderland Date: Mon, 12 Nov 2018 15:42:36 -0700 Subject: [PATCH] [MXNET-953] Fix oob memory read (#12631) --- src/operator/tensor/elemwise_unary_op_basic.cc | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/operator/tensor/elemwise_unary_op_basic.cc b/src/operator/tensor/elemwise_unary_op_basic.cc index a541085a6658..d1f1e08d95e3 100644 --- a/src/operator/tensor/elemwise_unary_op_basic.cc +++ b/src/operator/tensor/elemwise_unary_op_basic.cc @@ -492,7 +492,8 @@ void ShapeComputeCPU(const nnvm::NodeAttrs& attrs, CHECK_EQ(req.size(), 1U); const TBlob& in_data = inputs[0]; const TBlob& out_data = outputs[0]; - memcpy(out_data.dptr_, in_data.shape_.data(), in_data.ndim() * sizeof(int64_t)); + size_t type_size = mshadow::mshadow_sizeof(out_data.type_flag_); + memcpy(out_data.dptr_, in_data.shape_.data(), in_data.ndim() * type_size); } NNVM_REGISTER_OP(shape_array) @@ -542,8 +543,9 @@ void SizeComputeCPU(const nnvm::NodeAttrs& attrs, CHECK_EQ(req.size(), 1U); const TBlob& in_data = inputs[0]; const TBlob& out_data = outputs[0]; + size_t type_size = mshadow::mshadow_sizeof(out_data.type_flag_); const index_t size_var = in_data.Size(); - memcpy(out_data.dptr_, &size_var, 1U * sizeof(int64_t)); + memcpy(out_data.dptr_, &size_var, type_size); } NNVM_REGISTER_OP(size_array)