From 5d5865d377d3db9b936aa71936aab4176b3d3392 Mon Sep 17 00:00:00 2001 From: Kellen Sunderland Date: Fri, 21 Sep 2018 09:01:50 -0700 Subject: [PATCH] [MXNET-953] Fix oob memory read --- 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 c3e9c2dc91d0..dd0ad2b387dd 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)