Skip to content

Commit 5e03714

Browse files
authored
Fix pyre on builder.py (pytorch#6789)
Fix pyre on builder.py (pytorch#6789) Summary: Pull Request resolved: pytorch#6789 Reviewed By: lucylq Differential Revision: D65837880
1 parent 9cd57d6 commit 5e03714

File tree

1 file changed

+4
-4
lines changed

1 file changed

+4
-4
lines changed

extension/llm/export/builder.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@
88
# eager models, apply source transformations and quantization and export them to
99
# ExecuTorch.
1010

11+
# pyre-unsafe
12+
1113
import logging
1214
from enum import Enum
1315
from typing import Any, Callable, Dict, List, Optional
@@ -184,8 +186,6 @@ def export(self) -> "LLMEdgeManager":
184186
if hasattr(self.args, "qnn") and self.args.qnn:
185187
# TODO: this is temporary and export_for_training doesn't work with qnn either. We need a
186188
# functional graph. See issue https://github.com/pytorch/executorch/pull/4627 for more details
187-
# pyre-fixme[8]: Attribute has type `Optional[GraphModule]`; used as
188-
# `Module`.
189189
exported_module = torch.export.export(
190190
self.model,
191191
self.example_inputs,
@@ -194,14 +194,14 @@ def export(self) -> "LLMEdgeManager":
194194
strict=True,
195195
)
196196
else:
197-
# pyre-fixme[8]: Attribute has type `Optional[GraphModule]`; used as
198-
# `Module`.
199197
exported_module = export_for_training(
200198
self.model,
201199
self.example_inputs,
202200
kwargs=self.example_kwarg_inputs,
203201
dynamic_shapes=dynamic_shape,
204202
)
203+
# pyre-fixme[8]: Attribute has type `Optional[GraphModule]`; used as
204+
# `Module`.
205205
self.pre_autograd_graph_module = exported_module.module()
206206
if hasattr(self.args, "export_only") and self.args.export_only:
207207
torch.export.save(exported_module, self.args.output_name)

0 commit comments

Comments
 (0)