add SCATTER_API definition for scatter_mul in scatter.cpp & scatter.h (…

…#344) Co-authored-by: zenghongtai <[email protected]>
rusty1s · Dec 9, 2022 · fe83843 · fe83843
1 parent 111ffc4
commit fe83843
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/csrc/scatter.cpp b/csrc/scatter.cpp
@@ -239,9 +239,10 @@ scatter_sum(torch::Tensor src, torch::Tensor index, int64_t dim,
   return ScatterSum::apply(src, index, dim, optional_out, dim_size)[0];
 }
 
-torch::Tensor scatter_mul(torch::Tensor src, torch::Tensor index, int64_t dim,
-                          torch::optional<torch::Tensor> optional_out,
-                          torch::optional<int64_t> dim_size) {
+SCATTER_API torch::Tensor
+scatter_mul(torch::Tensor src, torch::Tensor index, int64_t dim,
+            torch::optional<torch::Tensor> optional_out,
+            torch::optional<int64_t> dim_size) {
   return ScatterMul::apply(src, index, dim, optional_out, dim_size)[0];
 }
 

diff --git a/csrc/scatter.h b/csrc/scatter.h
@@ -15,6 +15,11 @@ scatter_sum(torch::Tensor src, torch::Tensor index, int64_t dim,
             torch::optional<torch::Tensor> optional_out,
             torch::optional<int64_t> dim_size);
 
+SCATTER_API torch::Tensor
+scatter_mul(torch::Tensor src, torch::Tensor index, int64_t dim,
+            torch::optional<torch::Tensor> optional_out,
+            torch::optional<int64_t> dim_size);
+
 SCATTER_API torch::Tensor
 scatter_mean(torch::Tensor src, torch::Tensor index, int64_t dim,
              torch::optional<torch::Tensor> optional_out,