Update import of sdpa_with_kv_cache to custom_ops (#1470)

pytorch · Jan 22, 2025 · b2d8f2a · b2d8f2a
1 parent 45cd239
commit b2d8f2a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchchat/model.py b/torchchat/model.py
@@ -1025,7 +1025,7 @@ def apply_rotary_emb(x: Tensor, freqs_cis: Tensor) -> Tensor:
     # For quantized_decomposed ops
     from executorch.kernels import quantized  # no-qa
     # For llama::sdpa_with_kv_cache.out, preprocess ops
-    from executorch.extension.llm.custom_ops import sdpa_with_kv_cache  # no-qa
+    from executorch.extension.llm.custom_ops import custom_ops  # no-qa
 
     class PTEModel(nn.Module):
         def __init__(self, config, path) -> None: