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

diff --git a/caffe2/python/pybind_state_ideep.cc b/caffe2/python/pybind_state_ideep.cc
index 7057ca9..2b20af2 100644 (file)
--- 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<float>() && !ZeroDim(original_array)) {
+      if ((in_place && blob->IsType<itensor>())
+          || (meta.Match<float>() && !ZeroDim(original_array))) {
         FeedTensor(option, original_array, blob->GetMutable<itensor>());
       } else {
         DeviceOption cpu_option(option);
@@ -181,9 +183,9 @@ public:
         TensorFeeder<CPUContext> 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<Tensor>(new Tensor(