diff --git a/fbgemm_gpu/experimental/gen_ai/src/quantize/quantize.cpp b/fbgemm_gpu/experimental/gen_ai/src/quantize/quantize.cpp
index 2a88ee10f..a51674044 100644
--- a/fbgemm_gpu/experimental/gen_ai/src/quantize/quantize.cpp
+++ b/fbgemm_gpu/experimental/gen_ai/src/quantize/quantize.cpp
@@ -79,7 +79,7 @@ at::Tensor get_fp8_per_tensor_scale(
     c10::optional<at::Tensor> bs,
     c10::optional<at::Tensor> scale_ub); // scale upperbound
 
-TORCH_LIBRARY(fbgemm, m) {
+TORCH_LIBRARY_FRAGMENT(fbgemm, m) {
 #ifndef USE_ROCM
   // TODO: on AMD this throws "Undefined symbol" when loading
   // quantize_ops with