[Bugfix] Fix Machete unittests failing with NotImplementedError (vl…

…lm-project#9218) Signed-off-by: Sumit Dubey <sumit.dubey2@ibm.com>
sumitd2 · Nov 14, 2024 · 10a0e08 · 10a0e08
1 parent 5af3b76
commit 10a0e08
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/csrc/quantization/machete/machete_pytorch.cu b/csrc/quantization/machete/machete_pytorch.cu
@@ -89,6 +89,10 @@ torch::Tensor prepack_B(torch::Tensor const& B,
 TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
   m.impl("machete_prepack_B", &prepack_B);
   m.impl("machete_gemm", &gemm);
+}
+
+// use CatchAll since supported_schedules has no tensor arguments
+TORCH_LIBRARY_IMPL(TORCH_EXTENSION_NAME, CatchAll, m) {
   m.impl("machete_supported_schedules", &supported_schedules);
 }