From 89abb8d7118adb7e08bcdce43376a5c76a150c1c Mon Sep 17 00:00:00 2001
From: Jack-Khuu <jack.khuu.7@gmail.com>
Date: Tue, 21 Jan 2025 17:17:42 -0800
Subject: [PATCH] Fix import typo of spda

---
 torchchat/model.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/torchchat/model.py b/torchchat/model.py
index f50d2a8be..36cc74bdb 100644
--- 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: