Skip to content
Closed
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions extension/llm/export/builder.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,10 +29,10 @@

from executorch.extension.export_util.utils import export_to_edge, save_pte_program
from executorch.extension.llm.tokenizer.utils import get_tokenizer
from torch._export import capture_pre_autograd_graph
from torch.ao.quantization.quantize_pt2e import convert_pt2e, prepare_pt2e
from torch.ao.quantization.quantizer import Quantizer
from torch.ao.quantization.quantizer.composable_quantizer import ComposableQuantizer
from torch.export import export_for_training
from torch.nn.attention import SDPBackend

FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s"
Expand Down Expand Up @@ -190,9 +190,9 @@ def capture_pre_autograd_graph(self) -> "LLMEdgeManager":
strict=True,
).module()
else:
self.pre_autograd_graph_module = export_for_training(
self.pre_autograd_graph_module = capture_pre_autograd_graph(
self.model, self.example_inputs, dynamic_shapes=dynamic_shape
).module()
)

return self

Expand Down
Loading