quant_weights is a CPU function

torch-ext/torch_binding.cpp

@@ -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);
+  ops.impl("quant_weights", torch::kCPU, &symmetric_quantize_last_axis_of_tensor);
 }
 
 REGISTER_EXTENSION(TORCH_EXTENSION_NAME)