Skip to content

Commit 7d9f385

Browse files
author
Sanggyu Lee
committed
Add get_compile_config
1 parent 8f43620 commit 7d9f385

File tree

1 file changed

+4
-1
lines changed
  • test/modules/model/TinyLlamaWithFusedAttention

1 file changed

+4
-1
lines changed

test/modules/model/TinyLlamaWithFusedAttention/model.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import torch
2-
2+
from tico.config.v1 import CompileConfigV1
33
from tico.serialize.operators.adapters.llama_rmsnorm import patched_llama_rmsnorm
44
from tico.serialize.operators.adapters.onert.llama_attention import (
55
llama_attention_forward_adapter,
@@ -76,3 +76,6 @@ def get_example_inputs(self):
7676
generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
7777
print(generated_text)
7878
return captured_input, {}
79+
80+
def get_compile_config(self):
81+
return CompileConfigV1(convert_single_batch_lhs_const_bmm_to_fc=True)

0 commit comments

Comments
 (0)