Skip to content

GitLab

Unverified Commit 18511aed authored by Lucas Wilkinson's avatar Lucas Wilkinson Committed by GitHub
Browse files

[Bugfix] Fix Machete unittests failing with `NotImplementedError` (#9218)

parent 83ea5c72
Hide whitespace changes
Inline Side-by-side
Showing with 4 additions and 0 deletions
csrc/quantization/machete/machete_pytorch.cu
View file @ 18511aed
...@@ -89,6 +89,10 @@ torch::Tensor prepack_B(torch::Tensor const& B, ...@@ -89,6 +89,10 @@ torch::Tensor prepack_B(torch::Tensor const& B,
TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) { TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
m.impl("machete_prepack_B", &prepack_B); m.impl("machete_prepack_B", &prepack_B);
m.impl("machete_gemm", &gemm); 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); m.impl("machete_supported_schedules", &supported_schedules);
} }
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment