Align getitem with pytorch when index on different devices

oneflow/core/functional/impl/array_functor.cpp

@@ -1840,7 +1840,19 @@ class TensorGetItemFunctor {
 
     Shape shape(DimVector(target_dims.begin(), target_dims.end()));
     if (shape != *(result->shape())) { result = JUST(Reshape(result, shape)); }
-    if (!tensor_indices.empty()) { result = JUST(ApplyAdvancedIndexing(result, tensor_indices)); }
+    if (!tensor_indices.empty()) {
+      const auto x_device = JUST(x->device());
+      for (int64_t i = 0; i < tensor_indices.size(); ++i) {
+        const auto tensor_index = tensor_indices.at(i);
+        if (tensor_index == nullptr) { continue; }
+        const auto tensor_index_device = JUST(tensor_index->device());
+        if ((tensor_index_device->type() != x_device->type())
+            || (tensor_index_device->device_id() != x_device->device_id())) {
+          tensor_indices.at(i) = JUST(Copy(tensor_index, x_device->type(), x_device->device_id()));
+        }
+      }
+      result = JUST(ApplyAdvancedIndexing(result, tensor_indices));
+    }
 
     // TODO(): Returns a view of tensor `x`.
     if (result == x) { result = JUST(Identity(x)); }

python/oneflow/test/tensor/test_tensor_indexing.py

@@ -327,6 +327,22 @@ def test_list_indexing_using_scalar_tensor(test_case):
         for arg in GenArgList(arg_dict):
             arg[0](test_case, *arg[1:])
 
+    def test_indecies_on_different_devices_1_gpu(test_case):
+        x = flow.ones(3, 10)
+        y = flow.ones(3, 10, device=flow.device("cuda:0"))
+
+        x_idx = [flow.tensor([1, 2]), flow.tensor([2, 0], device=flow.device("cuda:0"))]
+        y_idx = [flow.tensor([1, 2], device=flow.device("cuda:0")), flow.tensor([2, 0])]
+
+        test_case.assertTrue(np.allclose(x[x_idx].numpy(), np.array([1, 1])))
+        test_case.assertTrue(np.allclose(y[y_idx].numpy(), np.array([1, 1])))
+
+    @flow.unittest.skip_unless_1n2d()
+    def test_indecies_on_different_devices_2_gpu(test_case):
+        x = flow.ones(3, 10, device=flow.device("cuda:0"))
+        idx = [flow.tensor([1, 2], device=flow.device("cuda:1")), flow.tensor([2, 0])]
+        test_case.assertTrue(np.allclose(x[idx].numpy(), np.array([1, 1])))
+
 
 if __name__ == "__main__":
     unittest.main()