diff --git a/backends/arm/test/models/test_llama.py b/backends/arm/test/models/test_llama.py index 973f62d2724..2656c12417d 100644 --- a/backends/arm/test/models/test_llama.py +++ b/backends/arm/test/models/test_llama.py @@ -52,7 +52,7 @@ def prepare_model(self): params_file, str ), "invalid input for --llama_inputs" else: - logging.warning( + logger.warning( "Skipping Llama test because of lack of input. To run use --llama_inputs <.pt> <.json>" ) return None, None, None @@ -61,6 +61,8 @@ def prepare_model(self): params_file ), "Invalid file paths" + logger.info("Running test_llama.py") + # TODO: Enable key value cache args = [ "--disable_dynamic_shape", @@ -112,9 +114,11 @@ def test_llama_tosa_MI(self): ) .export() .to_edge_transform_and_lower() - .check_count({"torch.ops.higher_order.executorch_call_delegate": 14}) + .check_count({"torch.ops.higher_order.executorch_call_delegate": 26}) .to_executorch() .run_method_and_compare_outputs( - inputs=llama_inputs, atol=1.8, rtol=0.01 # TODO: decrease tolerance + inputs=llama_inputs, + atol=4.3, + rtol=1.1, # TODO: MLETORCH-825 decrease tolerance ) )