Merge pull request #639 from take-cheeze/fix_take_dev

Workaround device boundary issue in Gather op
pfnet-research · Sep 4, 2019 · 6ef4b4a · 6ef4b4a
2 parents 4a31e71 + a00eb0c
commit 6ef4b4a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/runtime/ops/indexing.cc b/runtime/ops/indexing.cc
@@ -188,7 +188,7 @@ chainerx::Array GetItemGradOp::RunImpl(
 }
 
 chainerx::Array GatherOp::RunImpl(ChxVMState* st, const chainerx::Array& data, const chainerx::Array& indices) {
-    return data.Take(indices, axis);
+    return data.Take(indices.ToDevice(data.device()), axis);
 }
 
 chainerx::Array GatherElementsOp::RunImpl(ChxVMState* st, const chainerx::Array& data, const chainerx::Array& indices_) {