From 84f020fe098bfaad308bd820869fafa44d3d4fff Mon Sep 17 00:00:00 2001 From: "Gu, Jinghui" Date: Fri, 29 Mar 2019 14:06:09 -0700 Subject: [PATCH] Fix bug in tensor feed which caused crash due to wrong tensor type (#18552) Summary: In blob feeder for ideep device, the wrong device option is given and led to a crash issue. This patch aims to correct the device option to fix this bug. Pull Request resolved: https://github.com/pytorch/pytorch/pull/18552 Differential Revision: D14679838 Pulled By: yinghai fbshipit-source-id: bde11e6a6fe44822166881dcb7c9bd0b34b4ecf3 --- caffe2/python/pybind_state_ideep.cc | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/caffe2/python/pybind_state_ideep.cc b/caffe2/python/pybind_state_ideep.cc index 7057ca9..2b20af2 100644 --- a/caffe2/python/pybind_state_ideep.cc +++ b/caffe2/python/pybind_state_ideep.cc @@ -172,8 +172,10 @@ public: const auto npy_type = PyArray_TYPE(array); const TypeMeta &meta = NumpyTypeToCaffe(npy_type); + // TODO: if necessary, use dispatcher. - if (meta.Match() && !ZeroDim(original_array)) { + if ((in_place && blob->IsType()) + || (meta.Match() && !ZeroDim(original_array))) { FeedTensor(option, original_array, blob->GetMutable()); } else { DeviceOption cpu_option(option); @@ -181,9 +183,9 @@ public: TensorFeeder cpu_tensor_feeder; if (in_place) { cpu_tensor_feeder.FeedTensor( - option, + cpu_option, original_array, - BlobGetMutableTensor(blob, OptionToDevice(option).type()), + BlobGetMutableTensor(blob, OptionToDevice(cpu_option).type()), true); } else { blob->Reset(new Tensor( -- 2.7.4