fix bug

PaddlePaddle · Aug 21, 2022 · af52de9 · af52de9
1 parent 01dbf51
commit af52de9
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/cinn/hlir/op/contrib/scatter.cc b/cinn/hlir/op/contrib/scatter.cc
@@ -46,6 +46,15 @@ using common::CINNValuePack;
 
 ir::Tensor Scatter(
     const ir::Tensor &A, const ir::Tensor &B, const ir::Tensor &C, const Type &axis, const std::string &name) {
+  std::string extern_fun_name;
+  if (target.arch == common::Target::Arch::NVGPU) {
+    extern_fun_name.assign("cinn_host_find_int_from_start");
+  } else if (target.arch == common::Target::Arch::X86) {
+    extern_fun_name.assign("cinn_host_find_int_from_start");
+  } else {
+    LOG(FATAL) << "ScatterAssign only support X86 and NVGPU ! Please Check.\n";
+  }
+
   auto pos_axis = axis;
   if (pos_axis < 0) pos_axis += input->shape.size();
   std::vector<int> new_axes;