change TORCH_LIBRARY to TORCH_LIBRARY_FRAGMENT to prevent conflict be…

…tween cpu/mps
pytorch · Jan 30, 2025 · c08425b · c08425b
1 parent 7815262
commit c08425b
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/torchao/experimental/ops/linear_8bit_act_xbit_weight/op_linear_8bit_act_xbit_weight_aten.cpp b/torchao/experimental/ops/linear_8bit_act_xbit_weight/op_linear_8bit_act_xbit_weight_aten.cpp
@@ -60,7 +60,7 @@
       "_linear_8bit_act_" #weight_nbit "bit_weight",          \
       &linear_meta<weight_nbit, true>);
 
-TORCH_LIBRARY(torchao, m) {
+TORCH_LIBRARY_FRAGMENT(torchao, m) {
   DEFINE_OP(1);
   DEFINE_OP(2);
   DEFINE_OP(3);

diff --git a/torchao/experimental/ops/mps/linear_fp_act_xbit_weight_aten.mm b/torchao/experimental/ops/mps/linear_fp_act_xbit_weight_aten.mm
@@ -163,7 +163,7 @@ Tensor pack_weights_cpu_kernel(const Tensor& W) {
   return B;
 }
 
-TORCH_LIBRARY(torchao, m) {
+TORCH_LIBRARY_FRAGMENT(torchao, m) {
   m.def("_pack_weight_1bit(Tensor W) -> Tensor");
   m.def("_pack_weight_2bit(Tensor W) -> Tensor");
   m.def("_pack_weight_3bit(Tensor W) -> Tensor");