Skip to content

Commit 6fe74eb

Browse files
aorenstefacebook-github-bot
authored andcommitted
typing for decorators - fx/_compatibility (#2322)
Summary: X-link: ctrl-labs/src2#33884 X-link: pytorch/executorch#4810 Pull Request resolved: #2322 X-link: pytorch/pytorch#134054 See #131429 Reviewed By: laithsakka Differential Revision: D61493706 fbshipit-source-id: d2b3feeff2abf8610e4e9940a1b93b5f80777dc2
1 parent 85f994d commit 6fe74eb

File tree

1 file changed

+10
-3
lines changed

1 file changed

+10
-3
lines changed

torchrec/fx/tracer.py

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,8 @@
77

88
# pyre-strict
99

10-
from typing import Any, Callable, Dict, List, Optional, Union
10+
import typing
11+
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
1112

1213
import torch
1314
from torch.fx._compatibility import compatibility
@@ -109,7 +110,11 @@ def create_arg(self, a: Any) -> Argument:
109110
return self.create_node(
110111
"call_function",
111112
target=NoWait,
112-
args=self.create_arg((a._obj,)),
113+
# Ugh. This line seems to be triggering some bug in pyre - so
114+
# cast instead of fixme.
115+
args=typing.cast(
116+
Tuple[torch.fx.node.Argument, ...], self.create_arg((a._obj,))
117+
),
113118
kwargs={},
114119
type_expr=NoWait,
115120
)
@@ -159,4 +164,6 @@ def symbolic_trace(
159164
"""
160165
tracer = Tracer(leaf_modules)
161166
graph = tracer.trace(root, concrete_args)
162-
return torch.fx.GraphModule(root, graph)
167+
# Ugh. This line seems to be triggering some bug in pyre - so cast instead
168+
# of fixme.
169+
return torch.fx.GraphModule(typing.cast(torch.nn.Module, root), graph)

0 commit comments

Comments
 (0)