Fix wheel build CI jobs (#7372)

ghstack-source-id: f077c82325da4211cd16705e812e6ebed6713263 Pull Request resolved: #7367 Co-authored-by: Stephen Jia <[email protected]> (cherry picked from commit 18142f7)
pytorch · Dec 18, 2024 · e2277ac · e2277ac
1 parent 9069017
commit e2277ac
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/build/packaging/smoke_test.py b/build/packaging/smoke_test.py
@@ -17,7 +17,7 @@
 
 # Import custom ops. This requires portable_lib to be loaded first.
 from executorch.extension.llm.custom_ops import (  # noqa: F401, F403
-    sdpa_with_kv_cache,
+    custom_ops,
 )  # usort: skip
 
 # Import quantized ops. This requires portable_lib to be loaded first.