kernels-community
/

quantization-eetq

Model card Files Files and versions

danieldk HF Staff commited on Feb 17, 2025

Commit

8294a79

·

1 Parent(s): f1c3798

quant_weights is a CPU function

Files changed (1) hide show

torch-ext/torch_binding.cpp +1 -1

torch-ext/torch_binding.cpp CHANGED Viewed

@@ -13,7 +13,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.impl("preprocess_weights", torch::kCUDA, &preprocess_weights_cuda);
   ops.def("quant_weights(Tensor origin_weight, ScalarType quant_type,"
                       "bool return_unprocessed_quantized_tensor) -> Tensor[]");
-  ops.impl("quant_weights", torch::kCUDA, &symmetric_quantize_last_axis_of_tensor);
 }
 REGISTER_EXTENSION(TORCH_EXTENSION_NAME)

   ops.impl("preprocess_weights", torch::kCUDA, &preprocess_weights_cuda);
   ops.def("quant_weights(Tensor origin_weight, ScalarType quant_type,"
                       "bool return_unprocessed_quantized_tensor) -> Tensor[]");
+  ops.impl("quant_weights", torch::kCPU, &symmetric_quantize_last_axis_of_tensor);
 }
 REGISTER_EXTENSION(TORCH_EXTENSION_NAME)