@@ -126,9 +126,7 @@ def return_wrapper():
126126        trace_inputs_method  =  "get_upper_bound_inputs" 
127127        get_trace_inputs  =  get_inputs_adapter (
128128            (
129-                 # pyre-fixme[6]: For 1st argument expected `(...) -> Any` but got 
130-                 #  `Union[Module, Tensor]`. 
131-                 getattr (eager_module , trace_inputs_method )
129+                 getattr (eager_module , trace_inputs_method )  # type: ignore[arg-type] 
132130                if  hasattr (eager_module , trace_inputs_method )
133131                else  eager_module .get_random_inputs 
134132            ),
@@ -144,18 +142,14 @@ def return_wrapper():
144142        if  hasattr (eager_module , "get_dynamic_shapes" ):
145143            assert  capture_config  is  not None 
146144            assert  capture_config .enable_aot  is  True 
147-             # pyre-fixme[29]: `Union[nn.modules.module.Module, 
148-             #  torch._tensor.Tensor]` is not a function. 
149-             trace_dynamic_shapes  =  eager_module .get_dynamic_shapes ()
145+             trace_dynamic_shapes  =  eager_module .get_dynamic_shapes ()  # type: ignore[operator] 
150146            method_name_to_dynamic_shapes  =  {}
151147            for  method  in  methods :
152148                method_name_to_dynamic_shapes [method ] =  trace_dynamic_shapes 
153149
154150        memory_planning_pass  =  MemoryPlanningPass ()
155151        if  hasattr (eager_module , "get_memory_planning_pass" ):
156-             # pyre-fixme[29]: `Union[nn.modules.module.Module, 
157-             #  torch._tensor.Tensor]` is not a function. 
158-             memory_planning_pass  =  eager_module .get_memory_planning_pass ()
152+             memory_planning_pass  =  eager_module .get_memory_planning_pass ()  # type: ignore[operator] 
159153
160154        class  WrapperModule (nn .Module ):
161155            def  __init__ (self , method ):
@@ -172,7 +166,7 @@ def __init__(self, method):
172166                assert  method_name  ==  "forward" 
173167                ep  =  _export (
174168                    eager_module ,
175-                     method_input ,
169+                     method_input ,   # type: ignore[arg-type] 
176170                    dynamic_shapes = (
177171                        method_name_to_dynamic_shapes [method_name ]
178172                        if  method_name_to_dynamic_shapes 
@@ -184,7 +178,7 @@ def __init__(self, method):
184178            else :
185179                exported_methods [method_name ] =  export (
186180                    eager_module ,
187-                     method_input ,
181+                     method_input ,   # type: ignore[arg-type] 
188182                    dynamic_shapes = (
189183                        method_name_to_dynamic_shapes [method_name ]
190184                        if  method_name_to_dynamic_shapes 
@@ -220,9 +214,7 @@ def __init__(self, method):
220214
221215        # Get a function that creates random inputs appropriate for testing. 
222216        get_random_inputs_fn  =  get_inputs_adapter (
223-             # pyre-fixme[6]: For 1st argument expected `(...) -> Any` but got 
224-             #  `Union[Module, Tensor]`. 
225-             eager_module .get_random_inputs ,
217+             eager_module .get_random_inputs ,  # type: ignore[arg-type] 
226218            # all exported methods must have the same signature so just pick the first one. 
227219            methods [0 ],
228220        )
0 commit comments