From 18511aeda64b473314bb7727a97a220565e0af41 Mon Sep 17 00:00:00 2001
From: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com>
Date: Thu, 10 Oct 2024 13:39:56 -0400
Subject: [PATCH] [Bugfix] Fix Machete unittests failing with
 `NotImplementedError` (#9218)

---
 csrc/quantization/machete/machete_pytorch.cu | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/csrc/quantization/machete/machete_pytorch.cu b/csrc/quantization/machete/machete_pytorch.cu
index a27f1e7c83df9..ff037756f55ab 100644
--- 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);
 }