Fixes

Author: jackzhxng

examples/models/llama/runner/native.py

@@ -23,7 +23,7 @@
 from executorch.examples.models.llama.runner.generation import LlamaRunner
 
 # Note: import this after portable_lib
-# from executorch.extension.llm.custom_ops import sdpa_with_kv_cache  # noqa # usort: skip
+from executorch.extension.llm.custom_ops import sdpa_with_kv_cache  # noqa # usort: skip
 from executorch.kernels import quantized  # noqa
 
 
@@ -50,17 +50,11 @@ def forward(
         tokens: torch.Tensor,
         input_pos: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
-        # TODO: in LlamaRunner there is a generate function that automatically generates
-        # input_pos tensor and inputs it into the model. Atm TorchTune models use
-        # kwargs for the input_pos, so we will need to make some changes. At least
-        # for the time being, we can run the non-kv cache version of the Torchtune
-        # model with just the tokens like below.
-        return (self.model.forward((tokens,)))[0]
-        # return (
-        #     self.model.forward((tokens, input_pos))
-        #     if input_pos is not None
-        #     else self.model.forward((tokens,))
-        # )[0]
+        return (
+            self.model.forward((tokens, input_pos))
+            if input_pos is not None
+            else self.model.forward((tokens,))
+        )[0]
 
 
 def build_args_parser() -> argparse.ArgumentParser:
@@ -69,7 +63,7 @@ def build_args_parser() -> argparse.ArgumentParser:
 
     parser.add_argument(
         "--model",
-        default="llama",
+        default="llama3",
         choices=EXECUTORCH_DEFINED_MODELS + TORCHTUNE_DEFINED_MODELS,
     )
 

examples/models/llama3_2_vision/model.py

@@ -111,7 +111,7 @@ def __init__(self, **kwargs):
         # Load checkpoint.
         missing, unexpected = self.model_.load_state_dict(
             checkpoint,
-            strict=True,
+            strict=False,
             assign=True,
         )
         if kwargs.get("verbose", False):