From 5a0068cc6928539f96cb1a7409e4970142f55f1d Mon Sep 17 00:00:00 2001
From: Aaron Orenstein <aorenste@fb.com>
Date: Tue, 23 Jul 2024 14:48:25 -0700
Subject: [PATCH] [BE] mypy: disallow untyped decorators (#131428)

Untyped decorators strip the types from their decorated function so even if the underlying function is fully typed then callers to it don't get any benefit from type annotations.

Step 1 - Enable the error and override in all the offending files.

#131429

Pull Request resolved: https://github.com/pytorch/pytorch/pull/131428
Approved by: https://github.com/justinchuby, https://github.com/oulgen
---
 mypy.ini                                                         | 1 +
 test/test_torch.py                                               | 1 +
 torch/_decomp/decompositions.py                                  | 1 +
 torch/_decomp/decompositions_for_jvp.py                          | 1 +
 torch/_decomp/decompositions_for_rng.py                          | 1 +
 torch/_dynamo/convert_frame.py                                   | 1 +
 torch/_dynamo/output_graph.py                                    | 1 +
 torch/_dynamo/variables/script_object.py                         | 1 +
 torch/_export/__init__.py                                        | 1 +
 torch/_functorch/apis.py                                         | 1 +
 torch/_functorch/batch_norm_replacement.py                       | 1 +
 torch/_functorch/functional_call.py                              | 1 +
 torch/_higher_order_ops/auto_functionalize.py                    | 1 +
 torch/_higher_order_ops/effects.py                               | 1 +
 torch/_higher_order_ops/flex_attention.py                        | 1 +
 torch/_higher_order_ops/out_dtype.py                             | 1 +
 torch/_inductor/codecache.py                                     | 1 +
 torch/_inductor/compile_fx.py                                    | 1 +
 torch/_inductor/decomposition.py                                 | 1 +
 torch/_inductor/freezing.py                                      | 1 +
 torch/_inductor/fx_passes/post_grad.py                           | 1 +
 torch/_inductor/fx_passes/quantization.py                        | 1 +
 torch/_inductor/jagged_lowerings.py                              | 1 +
 torch/_inductor/kernel/conv.py                                   | 1 +
 torch/_inductor/lowering.py                                      | 1 +
 torch/_inductor/mkldnn_lowerings.py                              | 1 +
 torch/_inductor/pattern_matcher.py                               | 1 +
 torch/_inductor/runtime/triton_helpers.py                        | 1 +
 torch/_inductor/scheduler.py                                     | 1 +
 torch/_library/custom_ops.py                                     | 1 +
 torch/_library/infer_schema.py                                   | 1 +
 torch/_meta_registrations.py                                     | 1 +
 torch/_refs/__init__.py                                          | 1 +
 torch/_refs/_conversions.py                                      | 1 +
 torch/_refs/fft.py                                               | 1 +
 torch/_refs/linalg/__init__.py                                   | 1 +
 torch/_refs/nn/functional/__init__.py                            | 1 +
 torch/_refs/special/__init__.py                                  | 1 +
 torch/_subclasses/fake_tensor.py                                 | 1 +
 torch/ao/nn/quantized/dynamic/modules/rnn.py                     | 1 +
 torch/ao/nn/quantized/modules/embedding_ops.py                   | 1 +
 torch/ao/nn/quantized/modules/linear.py                          | 1 +
 torch/ao/nn/sparse/quantized/linear.py                           | 1 +
 torch/ao/ns/fx/utils.py                                          | 1 +
 torch/ao/quantization/experimental/adaround_fake_quantize.py     | 1 +
 torch/ao/quantization/fake_quantize.py                           | 1 +
 torch/ao/quantization/fx/_decomposed.py                          | 1 +
 torch/ao/quantization/observer.py                                | 1 +
 torch/ao/quantization/quantizer/xnnpack_quantizer_utils.py       | 1 +
 torch/backends/_nnapi/prepare.py                                 | 1 +
 torch/distributed/_composable/checkpoint_activation.py           | 1 +
 torch/distributed/_composable/fsdp/_fsdp_collectives.py          | 1 +
 torch/distributed/_composable/fsdp/_fsdp_state.py                | 1 +
 torch/distributed/_composable/fsdp/fully_shard.py                | 1 +
 torch/distributed/_composable/fully_shard.py                     | 1 +
 torch/distributed/_composable/replicate.py                       | 1 +
 torch/distributed/_spmd/experimental_ops.py                      | 1 +
 torch/distributed/_symmetric_memory/__init__.py                  | 1 +
 torch/distributed/_tensor/api.py                                 | 1 +
 torch/distributed/_tensor/ops/conv_ops.py                        | 1 +
 torch/distributed/_tensor/ops/embedding_ops.py                   | 1 +
 torch/distributed/_tensor/ops/experimental_ops.py                | 1 +
 torch/distributed/_tensor/ops/math_ops.py                        | 1 +
 torch/distributed/_tensor/ops/matrix_ops.py                      | 1 +
 torch/distributed/_tensor/ops/random_ops.py                      | 1 +
 torch/distributed/_tensor/ops/tensor_ops.py                      | 1 +
 torch/distributed/_tensor/ops/view_ops.py                        | 1 +
 torch/distributed/checkpoint/state_dict_loader.py                | 1 +
 torch/distributed/checkpoint/state_dict_saver.py                 | 1 +
 torch/distributed/optim/optimizer.py                             | 1 +
 torch/distributed/rpc/api.py                                     | 1 +
 torch/export/_trace.py                                           | 1 +
 torch/export/exported_program.py                                 | 1 +
 torch/fx/_lazy_graph_module.py                                   | 1 +
 torch/fx/_symbolic_trace.py                                      | 1 +
 torch/fx/experimental/graph_gradual_typechecker.py               | 1 +
 .../experimental/migrate_gradual_types/constraint_generator.py   | 1 +
 torch/fx/experimental/proxy_tensor.py                            | 1 +
 torch/fx/experimental/rewriter.py                                | 1 +
 torch/fx/graph.py                                                | 1 +
 torch/fx/graph_module.py                                         | 1 +
 torch/fx/interpreter.py                                          | 1 +
 torch/fx/node.py                                                 | 1 +
 torch/fx/operator_schemas.py                                     | 1 +
 torch/fx/passes/graph_manipulation.py                            | 1 +
 torch/fx/passes/infra/pass_manager.py                            | 1 +
 torch/fx/passes/operator_support.py                              | 1 +
 torch/fx/passes/param_fetch.py                                   | 1 +
 torch/fx/passes/runtime_assert.py                                | 1 +
 torch/fx/passes/split_module.py                                  | 1 +
 torch/fx/passes/split_utils.py                                   | 1 +
 torch/fx/passes/splitter_base.py                                 | 1 +
 torch/fx/passes/tools_common.py                                  | 1 +
 torch/fx/passes/utils/common.py                                  | 1 +
 torch/fx/passes/utils/fuser_utils.py                             | 1 +
 torch/fx/passes/utils/source_matcher_utils.py                    | 1 +
 torch/fx/subgraph_rewriter.py                                    | 1 +
 torch/fx/traceback.py                                            | 1 +
 torch/jit/_decompositions.py                                     | 1 +
 torch/masked/_ops.py                                             | 1 +
 torch/nn/attention/flex_attention.py                             | 1 +
 torch/nn/modules/container.py                                    | 1 +
 torch/nn/modules/rnn.py                                          | 1 +
 torch/nn/utils/clip_grad.py                                      | 1 +
 torch/nn/utils/parametrize.py                                    | 1 +
 torch/onnx/_internal/onnxruntime.py                              | 1 +
 torch/onnx/symbolic_caffe2.py                                    | 1 +
 torch/onnx/symbolic_helper.py                                    | 1 +
 torch/onnx/symbolic_opset10.py                                   | 1 +
 torch/onnx/symbolic_opset11.py                                   | 1 +
 torch/onnx/symbolic_opset12.py                                   | 1 +
 torch/onnx/symbolic_opset13.py                                   | 1 +
 torch/onnx/symbolic_opset14.py                                   | 1 +
 torch/onnx/symbolic_opset16.py                                   | 1 +
 torch/onnx/symbolic_opset17.py                                   | 1 +
 torch/onnx/symbolic_opset18.py                                   | 1 +
 torch/onnx/symbolic_opset20.py                                   | 1 +
 torch/onnx/symbolic_opset8.py                                    | 1 +
 torch/onnx/symbolic_opset9.py                                    | 1 +
 torch/optim/adadelta.py                                          | 1 +
 torch/optim/adam.py                                              | 1 +
 torch/optim/adamax.py                                            | 1 +
 torch/optim/adamw.py                                             | 1 +
 torch/optim/asgd.py                                              | 1 +
 torch/optim/nadam.py                                             | 1 +
 torch/optim/optimizer.py                                         | 1 +
 torch/optim/radam.py                                             | 1 +
 torch/optim/rmsprop.py                                           | 1 +
 torch/optim/rprop.py                                             | 1 +
 torch/signal/windows/windows.py                                  | 1 +
 torch/sparse/_triton_ops.py                                      | 1 +
 torch/testing/_internal/custom_op_db.py                          | 1 +
 torch/utils/_freeze.py                                           | 1 +
 torch/utils/checkpoint.py                                        | 1 +
 torch/utils/flop_counter.py                                      | 1 +
 135 files changed, 135 insertions(+)

diff --git a/mypy.ini b/mypy.ini
index c4fef0f5ba6..ebe6ea4d201 100644
--- a/mypy.ini
+++ b/mypy.ini
@@ -12,6 +12,7 @@ show_error_codes = True
 show_column_numbers = True
 check_untyped_defs = True
 disallow_untyped_defs = True
+disallow_untyped_decorators = True
 follow_imports = normal
 local_partial_types = True
 enable_error_code = possibly-undefined
diff --git a/test/test_torch.py b/test/test_torch.py
index 3e782eabe92..28ed40fd820 100644
--- a/test/test_torch.py
+++ b/test/test_torch.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 # Owner(s): ["module: tests"]
 
diff --git a/torch/_decomp/decompositions.py b/torch/_decomp/decompositions.py
index 174af5d2efd..251cbfb1f29 100644
--- a/torch/_decomp/decompositions.py
+++ b/torch/_decomp/decompositions.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import functools
 import itertools
diff --git a/torch/_decomp/decompositions_for_jvp.py b/torch/_decomp/decompositions_for_jvp.py
index ce47ac43d37..81e94110550 100644
--- a/torch/_decomp/decompositions_for_jvp.py
+++ b/torch/_decomp/decompositions_for_jvp.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import inspect
 from typing import Callable, Dict, List, Optional, Tuple
diff --git a/torch/_decomp/decompositions_for_rng.py b/torch/_decomp/decompositions_for_rng.py
index 74eb9b9240a..0a055b3e827 100644
--- a/torch/_decomp/decompositions_for_rng.py
+++ b/torch/_decomp/decompositions_for_rng.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import functools
 from collections import defaultdict
diff --git a/torch/_dynamo/convert_frame.py b/torch/_dynamo/convert_frame.py
index c7010513ab0..9657739c730 100644
--- a/torch/_dynamo/convert_frame.py
+++ b/torch/_dynamo/convert_frame.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 from __future__ import annotations
 
 import collections
diff --git a/torch/_dynamo/output_graph.py b/torch/_dynamo/output_graph.py
index 73e9daa4217..802cd98da7f 100644
--- a/torch/_dynamo/output_graph.py
+++ b/torch/_dynamo/output_graph.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import collections
 import contextlib
diff --git a/torch/_dynamo/variables/script_object.py b/torch/_dynamo/variables/script_object.py
index 92343719364..f563e37a78c 100644
--- a/torch/_dynamo/variables/script_object.py
+++ b/torch/_dynamo/variables/script_object.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import functools
 from typing import Dict
diff --git a/torch/_export/__init__.py b/torch/_export/__init__.py
index 49dda814bc7..11438a987ea 100644
--- a/torch/_export/__init__.py
+++ b/torch/_export/__init__.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import copy
 import dataclasses
diff --git a/torch/_functorch/apis.py b/torch/_functorch/apis.py
index 1b755550a8b..492e77a5c71 100644
--- a/torch/_functorch/apis.py
+++ b/torch/_functorch/apis.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 # NOTE: We allow Dynamo to see this file (via torch/_dynamo/trace_rules.py) so that it can
 #       trace through functorch transforms.
diff --git a/torch/_functorch/batch_norm_replacement.py b/torch/_functorch/batch_norm_replacement.py
index 672a8ce7695..90e4fec99b5 100644
--- a/torch/_functorch/batch_norm_replacement.py
+++ b/torch/_functorch/batch_norm_replacement.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import torch.nn as nn
 from torch._functorch.utils import exposed_in
diff --git a/torch/_functorch/functional_call.py b/torch/_functorch/functional_call.py
index 3aa781be930..2798dabef1e 100644
--- a/torch/_functorch/functional_call.py
+++ b/torch/_functorch/functional_call.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from collections import Counter
 from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
diff --git a/torch/_higher_order_ops/auto_functionalize.py b/torch/_higher_order_ops/auto_functionalize.py
index 0454ea12351..40178fa750f 100644
--- a/torch/_higher_order_ops/auto_functionalize.py
+++ b/torch/_higher_order_ops/auto_functionalize.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import warnings
 from typing import Any, Dict, List, Optional, Tuple, Union
diff --git a/torch/_higher_order_ops/effects.py b/torch/_higher_order_ops/effects.py
index 1da0764eee9..f20c87c7e58 100644
--- a/torch/_higher_order_ops/effects.py
+++ b/torch/_higher_order_ops/effects.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from enum import Enum
 from typing import Any, Dict, Optional, Tuple, Union
diff --git a/torch/_higher_order_ops/flex_attention.py b/torch/_higher_order_ops/flex_attention.py
index d273e79b880..f0db63aabd0 100644
--- a/torch/_higher_order_ops/flex_attention.py
+++ b/torch/_higher_order_ops/flex_attention.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from typing import Callable, Tuple, Union
 
diff --git a/torch/_higher_order_ops/out_dtype.py b/torch/_higher_order_ops/out_dtype.py
index 58bf02b8dfc..5c9ca4f3f16 100644
--- a/torch/_higher_order_ops/out_dtype.py
+++ b/torch/_higher_order_ops/out_dtype.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 
 import torch
diff --git a/torch/_inductor/codecache.py b/torch/_inductor/codecache.py
index a982928b6df..d8476b1d011 100644
--- a/torch/_inductor/codecache.py
+++ b/torch/_inductor/codecache.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from __future__ import annotations
 
diff --git a/torch/_inductor/compile_fx.py b/torch/_inductor/compile_fx.py
index 5195891c10b..a2e9a84f1f3 100644
--- a/torch/_inductor/compile_fx.py
+++ b/torch/_inductor/compile_fx.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import contextlib
 import functools
diff --git a/torch/_inductor/decomposition.py b/torch/_inductor/decomposition.py
index 0ae33145f3d..6ef6032fcf8 100644
--- a/torch/_inductor/decomposition.py
+++ b/torch/_inductor/decomposition.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import functools
 import logging
diff --git a/torch/_inductor/freezing.py b/torch/_inductor/freezing.py
index c62e7d4893e..f6fcde8f983 100644
--- a/torch/_inductor/freezing.py
+++ b/torch/_inductor/freezing.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from __future__ import annotations
 
diff --git a/torch/_inductor/fx_passes/post_grad.py b/torch/_inductor/fx_passes/post_grad.py
index df942842639..085bcf48d4f 100644
--- a/torch/_inductor/fx_passes/post_grad.py
+++ b/torch/_inductor/fx_passes/post_grad.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import functools
 import itertools
diff --git a/torch/_inductor/fx_passes/quantization.py b/torch/_inductor/fx_passes/quantization.py
index f40d5dc8c30..4aa547ab5d5 100644
--- a/torch/_inductor/fx_passes/quantization.py
+++ b/torch/_inductor/fx_passes/quantization.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import copy
 import functools
diff --git a/torch/_inductor/jagged_lowerings.py b/torch/_inductor/jagged_lowerings.py
index 7a6f1f4632b..c96c9f4ae2d 100644
--- a/torch/_inductor/jagged_lowerings.py
+++ b/torch/_inductor/jagged_lowerings.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from typing import List, Optional, Tuple, Union
 
diff --git a/torch/_inductor/kernel/conv.py b/torch/_inductor/kernel/conv.py
index 36e12777efb..21c91634bc7 100644
--- a/torch/_inductor/kernel/conv.py
+++ b/torch/_inductor/kernel/conv.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from __future__ import annotations
 
diff --git a/torch/_inductor/lowering.py b/torch/_inductor/lowering.py
index 23f2a94fbd2..2565bf9b19d 100644
--- a/torch/_inductor/lowering.py
+++ b/torch/_inductor/lowering.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import functools
 import itertools
diff --git a/torch/_inductor/mkldnn_lowerings.py b/torch/_inductor/mkldnn_lowerings.py
index 84ed0583fd1..8b7776f642f 100644
--- a/torch/_inductor/mkldnn_lowerings.py
+++ b/torch/_inductor/mkldnn_lowerings.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import functools
 from typing import List, Optional
diff --git a/torch/_inductor/pattern_matcher.py b/torch/_inductor/pattern_matcher.py
index 4427ac1352c..0ea342b6a93 100644
--- a/torch/_inductor/pattern_matcher.py
+++ b/torch/_inductor/pattern_matcher.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 """
 # Inductor Pattern Matcher
 
diff --git a/torch/_inductor/runtime/triton_helpers.py b/torch/_inductor/runtime/triton_helpers.py
index 453058ee130..a860e7fa853 100644
--- a/torch/_inductor/runtime/triton_helpers.py
+++ b/torch/_inductor/runtime/triton_helpers.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import triton
 import triton.language as tl
diff --git a/torch/_inductor/scheduler.py b/torch/_inductor/scheduler.py
index fb974b255d0..77d4091eba3 100644
--- a/torch/_inductor/scheduler.py
+++ b/torch/_inductor/scheduler.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: disallow-untyped-defs
 from __future__ import annotations
 
diff --git a/torch/_library/custom_ops.py b/torch/_library/custom_ops.py
index 57f813e52d2..12798d3c348 100644
--- a/torch/_library/custom_ops.py
+++ b/torch/_library/custom_ops.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import inspect
 import logging
diff --git a/torch/_library/infer_schema.py b/torch/_library/infer_schema.py
index 756fc150594..d901276ea18 100644
--- a/torch/_library/infer_schema.py
+++ b/torch/_library/infer_schema.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import inspect
 import typing
diff --git a/torch/_meta_registrations.py b/torch/_meta_registrations.py
index de53a2a42f8..5a4c804c168 100644
--- a/torch/_meta_registrations.py
+++ b/torch/_meta_registrations.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import math
 from enum import Enum
diff --git a/torch/_refs/__init__.py b/torch/_refs/__init__.py
index 55c072e777d..6491ebfc701 100644
--- a/torch/_refs/__init__.py
+++ b/torch/_refs/__init__.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import builtins
 import collections
diff --git a/torch/_refs/_conversions.py b/torch/_refs/_conversions.py
index b312f8f6ead..ecb5660a6a4 100644
--- a/torch/_refs/_conversions.py
+++ b/torch/_refs/_conversions.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import torch
 import torch._prims_common as utils
diff --git a/torch/_refs/fft.py b/torch/_refs/fft.py
index 06dda2d3acc..860b1a00222 100644
--- a/torch/_refs/fft.py
+++ b/torch/_refs/fft.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 import math
 
 from typing import Iterable, List, Literal, NamedTuple, Optional, Sequence, Tuple, Union
diff --git a/torch/_refs/linalg/__init__.py b/torch/_refs/linalg/__init__.py
index 411087b773e..322e91c1524 100644
--- a/torch/_refs/linalg/__init__.py
+++ b/torch/_refs/linalg/__init__.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from functools import partial
 
diff --git a/torch/_refs/nn/functional/__init__.py b/torch/_refs/nn/functional/__init__.py
index 01700aabca1..120dd86be24 100644
--- a/torch/_refs/nn/functional/__init__.py
+++ b/torch/_refs/nn/functional/__init__.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import math
 from functools import wraps
diff --git a/torch/_refs/special/__init__.py b/torch/_refs/special/__init__.py
index 1e98deaeb16..ab0641d4680 100644
--- a/torch/_refs/special/__init__.py
+++ b/torch/_refs/special/__init__.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import math
 from typing import Optional, Union
diff --git a/torch/_subclasses/fake_tensor.py b/torch/_subclasses/fake_tensor.py
index 1f28c60496e..55c4d20bb39 100644
--- a/torch/_subclasses/fake_tensor.py
+++ b/torch/_subclasses/fake_tensor.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 from __future__ import annotations
 
 import atexit
diff --git a/torch/ao/nn/quantized/dynamic/modules/rnn.py b/torch/ao/nn/quantized/dynamic/modules/rnn.py
index 9afab93d1a5..e85164a78f0 100644
--- a/torch/ao/nn/quantized/dynamic/modules/rnn.py
+++ b/torch/ao/nn/quantized/dynamic/modules/rnn.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import numbers
 import warnings
diff --git a/torch/ao/nn/quantized/modules/embedding_ops.py b/torch/ao/nn/quantized/modules/embedding_ops.py
index 7418aa38307..d8e3bec6411 100644
--- a/torch/ao/nn/quantized/modules/embedding_ops.py
+++ b/torch/ao/nn/quantized/modules/embedding_ops.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import torch
 import torch.nn as nn
diff --git a/torch/ao/nn/quantized/modules/linear.py b/torch/ao/nn/quantized/modules/linear.py
index 52b0a80a1c9..8f76ef2c81c 100644
--- a/torch/ao/nn/quantized/modules/linear.py
+++ b/torch/ao/nn/quantized/modules/linear.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from collections.abc import Iterable
 import torch
diff --git a/torch/ao/nn/sparse/quantized/linear.py b/torch/ao/nn/sparse/quantized/linear.py
index 26388e2e2c7..9c99eb949cb 100644
--- a/torch/ao/nn/sparse/quantized/linear.py
+++ b/torch/ao/nn/sparse/quantized/linear.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from typing import Optional
 
diff --git a/torch/ao/ns/fx/utils.py b/torch/ao/ns/fx/utils.py
index 16ac0c9c150..d3b9831136c 100644
--- a/torch/ao/ns/fx/utils.py
+++ b/torch/ao/ns/fx/utils.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import enum
 import operator
diff --git a/torch/ao/quantization/experimental/adaround_fake_quantize.py b/torch/ao/quantization/experimental/adaround_fake_quantize.py
index d035a02b047..232e497965d 100644
--- a/torch/ao/quantization/experimental/adaround_fake_quantize.py
+++ b/torch/ao/quantization/experimental/adaround_fake_quantize.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from typing import Tuple
 
diff --git a/torch/ao/quantization/fake_quantize.py b/torch/ao/quantization/fake_quantize.py
index b921df39217..da4880a471f 100644
--- a/torch/ao/quantization/fake_quantize.py
+++ b/torch/ao/quantization/fake_quantize.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 """Implements modules  used to perform fake quantization."""
 
diff --git a/torch/ao/quantization/fx/_decomposed.py b/torch/ao/quantization/fx/_decomposed.py
index 72ce4b2471f..1d48af7ecf7 100644
--- a/torch/ao/quantization/fx/_decomposed.py
+++ b/torch/ao/quantization/fx/_decomposed.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import math
 from typing import Optional, Tuple
diff --git a/torch/ao/quantization/observer.py b/torch/ao/quantization/observer.py
index eef0e54cfad..3fbd2af046a 100644
--- a/torch/ao/quantization/observer.py
+++ b/torch/ao/quantization/observer.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 """
 This module implements observers which are used to collect statistics about
diff --git a/torch/ao/quantization/quantizer/xnnpack_quantizer_utils.py b/torch/ao/quantization/quantizer/xnnpack_quantizer_utils.py
index 928ee0d3ac4..4dbc5ad15a8 100644
--- a/torch/ao/quantization/quantizer/xnnpack_quantizer_utils.py
+++ b/torch/ao/quantization/quantizer/xnnpack_quantizer_utils.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import itertools
 import operator
diff --git a/torch/backends/_nnapi/prepare.py b/torch/backends/_nnapi/prepare.py
index 6ba389902c9..e78154bd7eb 100644
--- a/torch/backends/_nnapi/prepare.py
+++ b/torch/backends/_nnapi/prepare.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from typing import List, Optional
 
diff --git a/torch/distributed/_composable/checkpoint_activation.py b/torch/distributed/_composable/checkpoint_activation.py
index f1464d9e804..fcee2a57a07 100644
--- a/torch/distributed/_composable/checkpoint_activation.py
+++ b/torch/distributed/_composable/checkpoint_activation.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from contextlib import contextmanager, nullcontext
 from typing import Any, ContextManager, Dict, Optional, Tuple
diff --git a/torch/distributed/_composable/fsdp/_fsdp_collectives.py b/torch/distributed/_composable/fsdp/_fsdp_collectives.py
index 3d8642e8a24..e9444848c4d 100644
--- a/torch/distributed/_composable/fsdp/_fsdp_collectives.py
+++ b/torch/distributed/_composable/fsdp/_fsdp_collectives.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 from typing import List, NamedTuple, Optional, Tuple, Union
 
 import torch
diff --git a/torch/distributed/_composable/fsdp/_fsdp_state.py b/torch/distributed/_composable/fsdp/_fsdp_state.py
index f32b3c18ecd..3aad8e2226c 100644
--- a/torch/distributed/_composable/fsdp/_fsdp_state.py
+++ b/torch/distributed/_composable/fsdp/_fsdp_state.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import functools
 import logging
diff --git a/torch/distributed/_composable/fsdp/fully_shard.py b/torch/distributed/_composable/fsdp/fully_shard.py
index 5f686dbf6c9..795fd983f78 100644
--- a/torch/distributed/_composable/fsdp/fully_shard.py
+++ b/torch/distributed/_composable/fsdp/fully_shard.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import functools
 from typing import Any, cast, Iterable, List, NoReturn, Optional, Union
diff --git a/torch/distributed/_composable/fully_shard.py b/torch/distributed/_composable/fully_shard.py
index 06b121aef80..4afa0f43107 100644
--- a/torch/distributed/_composable/fully_shard.py
+++ b/torch/distributed/_composable/fully_shard.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 from typing import Callable, Iterable, Optional, Union
 from typing_extensions import deprecated
 
diff --git a/torch/distributed/_composable/replicate.py b/torch/distributed/_composable/replicate.py
index 391022a5805..2aa5d0fc3d3 100644
--- a/torch/distributed/_composable/replicate.py
+++ b/torch/distributed/_composable/replicate.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import weakref
 from typing import Any, cast, Dict, Iterable, List, NoReturn, Optional, Set, Tuple
diff --git a/torch/distributed/_spmd/experimental_ops.py b/torch/distributed/_spmd/experimental_ops.py
index 7039822c41c..f8c8f8804c5 100644
--- a/torch/distributed/_spmd/experimental_ops.py
+++ b/torch/distributed/_spmd/experimental_ops.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 # Copyright (c) Meta Platforms, Inc. and affiliates
 from typing import cast, List, Optional, Sequence, Tuple
diff --git a/torch/distributed/_symmetric_memory/__init__.py b/torch/distributed/_symmetric_memory/__init__.py
index eafcef0228e..2f7d861514b 100644
--- a/torch/distributed/_symmetric_memory/__init__.py
+++ b/torch/distributed/_symmetric_memory/__init__.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 import socket
 import uuid
 
diff --git a/torch/distributed/_tensor/api.py b/torch/distributed/_tensor/api.py
index d8c12b8cf35..fe889703d1a 100644
--- a/torch/distributed/_tensor/api.py
+++ b/torch/distributed/_tensor/api.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 # Copyright (c) Meta Platforms, Inc. and affiliates
 import inspect
diff --git a/torch/distributed/_tensor/ops/conv_ops.py b/torch/distributed/_tensor/ops/conv_ops.py
index 24e75593064..3a3b743dc47 100644
--- a/torch/distributed/_tensor/ops/conv_ops.py
+++ b/torch/distributed/_tensor/ops/conv_ops.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # Copyright (c) Meta Platforms, Inc. and affiliates
 # implement matrix related ops for distributed tensor
 from typing import List
diff --git a/torch/distributed/_tensor/ops/embedding_ops.py b/torch/distributed/_tensor/ops/embedding_ops.py
index 9e3e9e5eb4c..a374022ecf9 100644
--- a/torch/distributed/_tensor/ops/embedding_ops.py
+++ b/torch/distributed/_tensor/ops/embedding_ops.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 # Copyright (c) Meta Platforms, Inc. and affiliates
 # implement matrix related ops for distributed tensor
diff --git a/torch/distributed/_tensor/ops/experimental_ops.py b/torch/distributed/_tensor/ops/experimental_ops.py
index 0d2ac34b8bc..1edfb8d294f 100644
--- a/torch/distributed/_tensor/ops/experimental_ops.py
+++ b/torch/distributed/_tensor/ops/experimental_ops.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # Copyright (c) Meta Platforms, Inc. and affiliates
 # implement matrix related ops for distributed tensor
 
diff --git a/torch/distributed/_tensor/ops/math_ops.py b/torch/distributed/_tensor/ops/math_ops.py
index 15d06b1fe6e..662c1619e2f 100644
--- a/torch/distributed/_tensor/ops/math_ops.py
+++ b/torch/distributed/_tensor/ops/math_ops.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 # Copyright (c) Meta Platforms, Inc. and affiliates
 import math
diff --git a/torch/distributed/_tensor/ops/matrix_ops.py b/torch/distributed/_tensor/ops/matrix_ops.py
index 802b6b8e7dd..2815d14d949 100644
--- a/torch/distributed/_tensor/ops/matrix_ops.py
+++ b/torch/distributed/_tensor/ops/matrix_ops.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # Copyright (c) Meta Platforms, Inc. and affiliates
 # implement matrix related ops for distributed tensor
 
diff --git a/torch/distributed/_tensor/ops/random_ops.py b/torch/distributed/_tensor/ops/random_ops.py
index e968216744f..f54cfea7b00 100644
--- a/torch/distributed/_tensor/ops/random_ops.py
+++ b/torch/distributed/_tensor/ops/random_ops.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # Copyright (c) Meta Platforms, Inc. and affiliates
 import torch
 from torch.distributed._tensor._op_schema import (
diff --git a/torch/distributed/_tensor/ops/tensor_ops.py b/torch/distributed/_tensor/ops/tensor_ops.py
index e8a324bb400..e87fbb53c7c 100644
--- a/torch/distributed/_tensor/ops/tensor_ops.py
+++ b/torch/distributed/_tensor/ops/tensor_ops.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 # Copyright (c) Meta Platforms, Inc. and affiliates
 from typing import cast, List, Optional, Sequence, Tuple
diff --git a/torch/distributed/_tensor/ops/view_ops.py b/torch/distributed/_tensor/ops/view_ops.py
index ea088b7377a..23e303976ff 100644
--- a/torch/distributed/_tensor/ops/view_ops.py
+++ b/torch/distributed/_tensor/ops/view_ops.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 # Copyright (c) Meta Platforms, Inc. and affiliates
 from dataclasses import dataclass
diff --git a/torch/distributed/checkpoint/state_dict_loader.py b/torch/distributed/checkpoint/state_dict_loader.py
index c4d1d853e9c..1d92327e9e4 100644
--- a/torch/distributed/checkpoint/state_dict_loader.py
+++ b/torch/distributed/checkpoint/state_dict_loader.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import os
 import warnings
diff --git a/torch/distributed/checkpoint/state_dict_saver.py b/torch/distributed/checkpoint/state_dict_saver.py
index 20abc2212f5..307c6d8d4a9 100644
--- a/torch/distributed/checkpoint/state_dict_saver.py
+++ b/torch/distributed/checkpoint/state_dict_saver.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import inspect
 import os
diff --git a/torch/distributed/optim/optimizer.py b/torch/distributed/optim/optimizer.py
index 65df14770c2..f52fe488c46 100644
--- a/torch/distributed/optim/optimizer.py
+++ b/torch/distributed/optim/optimizer.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import logging
 from collections import defaultdict
diff --git a/torch/distributed/rpc/api.py b/torch/distributed/rpc/api.py
index 5fc9e61aa55..fbe77e46574 100644
--- a/torch/distributed/rpc/api.py
+++ b/torch/distributed/rpc/api.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 
 import collections
diff --git a/torch/export/_trace.py b/torch/export/_trace.py
index 3e2511433ad..b30b24cd1b7 100644
--- a/torch/export/_trace.py
+++ b/torch/export/_trace.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import dataclasses
 import functools
diff --git a/torch/export/exported_program.py b/torch/export/exported_program.py
index 1d7b77e5aaa..c7a38b9bb5c 100644
--- a/torch/export/exported_program.py
+++ b/torch/export/exported_program.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import contextlib
 import copy
diff --git a/torch/fx/_lazy_graph_module.py b/torch/fx/_lazy_graph_module.py
index 79a18de12f3..dc954b21b44 100644
--- a/torch/fx/_lazy_graph_module.py
+++ b/torch/fx/_lazy_graph_module.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from contextlib import contextmanager
 
diff --git a/torch/fx/_symbolic_trace.py b/torch/fx/_symbolic_trace.py
index 25a342f064c..bd8d4f42669 100644
--- a/torch/fx/_symbolic_trace.py
+++ b/torch/fx/_symbolic_trace.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import builtins
 import copy
diff --git a/torch/fx/experimental/graph_gradual_typechecker.py b/torch/fx/experimental/graph_gradual_typechecker.py
index a6ac80fd72f..fb49795a06f 100644
--- a/torch/fx/experimental/graph_gradual_typechecker.py
+++ b/torch/fx/experimental/graph_gradual_typechecker.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from functools import reduce
 import torch
diff --git a/torch/fx/experimental/migrate_gradual_types/constraint_generator.py b/torch/fx/experimental/migrate_gradual_types/constraint_generator.py
index e04fc26b408..952dde662f2 100644
--- a/torch/fx/experimental/migrate_gradual_types/constraint_generator.py
+++ b/torch/fx/experimental/migrate_gradual_types/constraint_generator.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import torch
 import operator
diff --git a/torch/fx/experimental/proxy_tensor.py b/torch/fx/experimental/proxy_tensor.py
index 8fd528360ce..540f95a3d57 100644
--- a/torch/fx/experimental/proxy_tensor.py
+++ b/torch/fx/experimental/proxy_tensor.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # Copyright (c) Facebook, Inc. and its affiliates.
 # All rights reserved.
 #
diff --git a/torch/fx/experimental/rewriter.py b/torch/fx/experimental/rewriter.py
index e0e3377b7b4..3647ca59153 100644
--- a/torch/fx/experimental/rewriter.py
+++ b/torch/fx/experimental/rewriter.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import ast
 import inspect
diff --git a/torch/fx/graph.py b/torch/fx/graph.py
index 01e91182c82..60c864b43b7 100644
--- a/torch/fx/graph.py
+++ b/torch/fx/graph.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from collections import defaultdict
 from .node import Node, Argument, Target, map_arg, _type_repr, _get_qualified_name
diff --git a/torch/fx/graph_module.py b/torch/fx/graph_module.py
index 90c8ca3cf18..9cbf6b7e9db 100644
--- a/torch/fx/graph_module.py
+++ b/torch/fx/graph_module.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import contextlib
 import copy
diff --git a/torch/fx/interpreter.py b/torch/fx/interpreter.py
index ebe4f3cbfe9..5f7dfa28b57 100644
--- a/torch/fx/interpreter.py
+++ b/torch/fx/interpreter.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from .graph_module import GraphModule
 from ._lazy_graph_module import _make_graph_module
diff --git a/torch/fx/node.py b/torch/fx/node.py
index 67c2eaebb89..987c64769ec 100644
--- a/torch/fx/node.py
+++ b/torch/fx/node.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # Nodes represent a definition of a value in our graph of operators.
 from typing import TYPE_CHECKING, Union, Callable, Any, Tuple, List, Optional, Dict, Set
 from ._compatibility import compatibility
diff --git a/torch/fx/operator_schemas.py b/torch/fx/operator_schemas.py
index 04be7d139da..68fb122b9a8 100644
--- a/torch/fx/operator_schemas.py
+++ b/torch/fx/operator_schemas.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import torch
 import inspect
diff --git a/torch/fx/passes/graph_manipulation.py b/torch/fx/passes/graph_manipulation.py
index 36c59cb31af..4f177046ea6 100644
--- a/torch/fx/passes/graph_manipulation.py
+++ b/torch/fx/passes/graph_manipulation.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from typing import Any, Dict, List, NamedTuple, Optional
 
diff --git a/torch/fx/passes/infra/pass_manager.py b/torch/fx/passes/infra/pass_manager.py
index 29540fa447e..9624fd5c8a4 100644
--- a/torch/fx/passes/infra/pass_manager.py
+++ b/torch/fx/passes/infra/pass_manager.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import inspect
 import logging
diff --git a/torch/fx/passes/operator_support.py b/torch/fx/passes/operator_support.py
index 57edabc0a55..d4133e7eb9a 100644
--- a/torch/fx/passes/operator_support.py
+++ b/torch/fx/passes/operator_support.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import abc
 import typing as t
diff --git a/torch/fx/passes/param_fetch.py b/torch/fx/passes/param_fetch.py
index 5979e29fcc6..2e32eda89a2 100644
--- a/torch/fx/passes/param_fetch.py
+++ b/torch/fx/passes/param_fetch.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 from torch.fx.graph_module import GraphModule
 from typing import Any, Callable, Dict, List, Tuple, Type
 import torch
diff --git a/torch/fx/passes/runtime_assert.py b/torch/fx/passes/runtime_assert.py
index 72b49e225aa..6e12ff52b80 100644
--- a/torch/fx/passes/runtime_assert.py
+++ b/torch/fx/passes/runtime_assert.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import logging
 import operator
diff --git a/torch/fx/passes/split_module.py b/torch/fx/passes/split_module.py
index 7682129f7ab..fba516d74be 100644
--- a/torch/fx/passes/split_module.py
+++ b/torch/fx/passes/split_module.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import inspect
 from typing import Any, Callable, Dict, List, Optional, Set
diff --git a/torch/fx/passes/split_utils.py b/torch/fx/passes/split_utils.py
index 38aa56064db..d8254bd474b 100644
--- a/torch/fx/passes/split_utils.py
+++ b/torch/fx/passes/split_utils.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import copy
 from dataclasses import dataclass, field
diff --git a/torch/fx/passes/splitter_base.py b/torch/fx/passes/splitter_base.py
index f4aa439b409..53ad10309f4 100644
--- a/torch/fx/passes/splitter_base.py
+++ b/torch/fx/passes/splitter_base.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import argparse
 import copy
diff --git a/torch/fx/passes/tools_common.py b/torch/fx/passes/tools_common.py
index aac071ace8c..6689af96e13 100644
--- a/torch/fx/passes/tools_common.py
+++ b/torch/fx/passes/tools_common.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from typing import List, Tuple, Union, Dict, Any, Set, Mapping, Optional
 import collections
diff --git a/torch/fx/passes/utils/common.py b/torch/fx/passes/utils/common.py
index ba2ae45aabf..95eadd58499 100644
--- a/torch/fx/passes/utils/common.py
+++ b/torch/fx/passes/utils/common.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from typing import Dict, Tuple
 
diff --git a/torch/fx/passes/utils/fuser_utils.py b/torch/fx/passes/utils/fuser_utils.py
index 324e8a67801..007a79e708f 100644
--- a/torch/fx/passes/utils/fuser_utils.py
+++ b/torch/fx/passes/utils/fuser_utils.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import copy
 from queue import SimpleQueue
diff --git a/torch/fx/passes/utils/source_matcher_utils.py b/torch/fx/passes/utils/source_matcher_utils.py
index 0f2650ea8d4..82596f4d16d 100644
--- a/torch/fx/passes/utils/source_matcher_utils.py
+++ b/torch/fx/passes/utils/source_matcher_utils.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from dataclasses import dataclass, field
 from torch.fx.graph import Graph
diff --git a/torch/fx/subgraph_rewriter.py b/torch/fx/subgraph_rewriter.py
index 3106daca0b1..419337a1768 100644
--- a/torch/fx/subgraph_rewriter.py
+++ b/torch/fx/subgraph_rewriter.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 from .graph_module import GraphModule
 from .graph import Graph
 from .node import Node
diff --git a/torch/fx/traceback.py b/torch/fx/traceback.py
index 4e72a8011f6..77aee26c3fc 100644
--- a/torch/fx/traceback.py
+++ b/torch/fx/traceback.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import traceback
 from contextlib import contextmanager
diff --git a/torch/jit/_decompositions.py b/torch/jit/_decompositions.py
index 8ac456be482..b4cb8194d26 100644
--- a/torch/jit/_decompositions.py
+++ b/torch/jit/_decompositions.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import torch
 from torch import Tensor
diff --git a/torch/masked/_ops.py b/torch/masked/_ops.py
index 26094459c17..7409c9cf938 100644
--- a/torch/masked/_ops.py
+++ b/torch/masked/_ops.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import warnings
 from typing import Any, List, Optional, Tuple, TYPE_CHECKING, Union
diff --git a/torch/nn/attention/flex_attention.py b/torch/nn/attention/flex_attention.py
index 3d3d521ce92..e56603b05c1 100644
--- a/torch/nn/attention/flex_attention.py
+++ b/torch/nn/attention/flex_attention.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 # flake8: noqa C101
 """This module implements the user facing API for flex_attention in PyTorch."""
diff --git a/torch/nn/modules/container.py b/torch/nn/modules/container.py
index 6ce4f2bd2c9..30992e394b0 100644
--- a/torch/nn/modules/container.py
+++ b/torch/nn/modules/container.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import operator
 from collections import abc as container_abcs, OrderedDict
diff --git a/torch/nn/modules/rnn.py b/torch/nn/modules/rnn.py
index 2870eee2941..8e3c6c0b357 100644
--- a/torch/nn/modules/rnn.py
+++ b/torch/nn/modules/rnn.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import math
 import numbers
diff --git a/torch/nn/utils/clip_grad.py b/torch/nn/utils/clip_grad.py
index 327b7809396..1d7834e53b6 100644
--- a/torch/nn/utils/clip_grad.py
+++ b/torch/nn/utils/clip_grad.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import functools
 from typing import cast, Dict, Iterable, List, Optional, Tuple, Union
diff --git a/torch/nn/utils/parametrize.py b/torch/nn/utils/parametrize.py
index 6a66ea58699..d4946604bcb 100644
--- a/torch/nn/utils/parametrize.py
+++ b/torch/nn/utils/parametrize.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import collections
 import copyreg
diff --git a/torch/onnx/_internal/onnxruntime.py b/torch/onnx/_internal/onnxruntime.py
index 4024e51b14f..b9d6bce1651 100644
--- a/torch/onnx/_internal/onnxruntime.py
+++ b/torch/onnx/_internal/onnxruntime.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import dataclasses
 import importlib
diff --git a/torch/onnx/symbolic_caffe2.py b/torch/onnx/symbolic_caffe2.py
index 83a2ff6c32e..1fa38489e6b 100644
--- a/torch/onnx/symbolic_caffe2.py
+++ b/torch/onnx/symbolic_caffe2.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 # mypy: disable-error-code=arg-type
 import importlib
diff --git a/torch/onnx/symbolic_helper.py b/torch/onnx/symbolic_helper.py
index 395dd4febfa..ad2552ee9a9 100644
--- a/torch/onnx/symbolic_helper.py
+++ b/torch/onnx/symbolic_helper.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from __future__ import annotations
 
diff --git a/torch/onnx/symbolic_opset10.py b/torch/onnx/symbolic_opset10.py
index c30452185c5..48fac3b1af6 100644
--- a/torch/onnx/symbolic_opset10.py
+++ b/torch/onnx/symbolic_opset10.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 # mypy: disable-error-code=arg-type
 from __future__ import annotations
diff --git a/torch/onnx/symbolic_opset11.py b/torch/onnx/symbolic_opset11.py
index eafd440639e..8f1f1379f77 100644
--- a/torch/onnx/symbolic_opset11.py
+++ b/torch/onnx/symbolic_opset11.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 # mypy: disable-error-code=arg-type
 """This file exports ONNX ops for opset 11."""
diff --git a/torch/onnx/symbolic_opset12.py b/torch/onnx/symbolic_opset12.py
index 7aaefd37201..296fc03d585 100644
--- a/torch/onnx/symbolic_opset12.py
+++ b/torch/onnx/symbolic_opset12.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 # mypy: disable-error-code=arg-type
 from __future__ import annotations
diff --git a/torch/onnx/symbolic_opset13.py b/torch/onnx/symbolic_opset13.py
index e31416ae2bc..f92b15c0d1d 100644
--- a/torch/onnx/symbolic_opset13.py
+++ b/torch/onnx/symbolic_opset13.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 # EDITING THIS FILE? READ THIS FIRST!
 # see Note [Edit Symbolic Files] in README.md
diff --git a/torch/onnx/symbolic_opset14.py b/torch/onnx/symbolic_opset14.py
index b7bf47192a2..02e9f374ab7 100644
--- a/torch/onnx/symbolic_opset14.py
+++ b/torch/onnx/symbolic_opset14.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 # mypy: disable-error-code=arg-type
 """This file exports ONNX ops for opset 14.
diff --git a/torch/onnx/symbolic_opset16.py b/torch/onnx/symbolic_opset16.py
index 696f1fecb02..1b8388ee5c9 100644
--- a/torch/onnx/symbolic_opset16.py
+++ b/torch/onnx/symbolic_opset16.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 """This file exports ONNX ops for opset 16.
 
diff --git a/torch/onnx/symbolic_opset17.py b/torch/onnx/symbolic_opset17.py
index 884dcf6f4bb..4e8f66ea08b 100644
--- a/torch/onnx/symbolic_opset17.py
+++ b/torch/onnx/symbolic_opset17.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 # mypy: disable-error-code=arg-type
 """This file exports ONNX ops for opset 17.
diff --git a/torch/onnx/symbolic_opset18.py b/torch/onnx/symbolic_opset18.py
index ee441b59ff0..2ce0e882c23 100644
--- a/torch/onnx/symbolic_opset18.py
+++ b/torch/onnx/symbolic_opset18.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 """This file exports ONNX ops for opset 18.
 
diff --git a/torch/onnx/symbolic_opset20.py b/torch/onnx/symbolic_opset20.py
index fe4fa1e6231..0eb58e42e6f 100644
--- a/torch/onnx/symbolic_opset20.py
+++ b/torch/onnx/symbolic_opset20.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 """This file exports ONNX ops for opset 20.
 
diff --git a/torch/onnx/symbolic_opset8.py b/torch/onnx/symbolic_opset8.py
index 191a3259efb..1e7d868fe5d 100644
--- a/torch/onnx/symbolic_opset8.py
+++ b/torch/onnx/symbolic_opset8.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 """
 Note [ONNX operators that are added/updated from opset 8 to opset 9]
diff --git a/torch/onnx/symbolic_opset9.py b/torch/onnx/symbolic_opset9.py
index c7e7f231048..5b3acb28348 100644
--- a/torch/onnx/symbolic_opset9.py
+++ b/torch/onnx/symbolic_opset9.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 # mypy: disable-error-code=arg-type
 """This file exports ONNX ops for opset 9.
diff --git a/torch/optim/adadelta.py b/torch/optim/adadelta.py
index f1391df06e2..40e3c82ec0a 100644
--- a/torch/optim/adadelta.py
+++ b/torch/optim/adadelta.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from typing import Any, Dict, List, Optional, Union
 
diff --git a/torch/optim/adam.py b/torch/optim/adam.py
index 8d8a9cdca35..fb1978ddd8a 100644
--- a/torch/optim/adam.py
+++ b/torch/optim/adam.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from typing import List, Optional, Tuple, Union
 
diff --git a/torch/optim/adamax.py b/torch/optim/adamax.py
index 4c5cc807057..d3eaca53df6 100644
--- a/torch/optim/adamax.py
+++ b/torch/optim/adamax.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from typing import List, Optional, Tuple, Union
 
diff --git a/torch/optim/adamw.py b/torch/optim/adamw.py
index c117f45b5e0..b1273545c98 100644
--- a/torch/optim/adamw.py
+++ b/torch/optim/adamw.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from typing import cast, List, Optional, Tuple, Union
 
diff --git a/torch/optim/asgd.py b/torch/optim/asgd.py
index a7c89d6b922..eafa1a6adee 100644
--- a/torch/optim/asgd.py
+++ b/torch/optim/asgd.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from typing import List, Optional, Tuple, Union
 
diff --git a/torch/optim/nadam.py b/torch/optim/nadam.py
index 251898ceec7..6aa626510d9 100644
--- a/torch/optim/nadam.py
+++ b/torch/optim/nadam.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 r"""Implementation for the NAdam algorithm."""
 from typing import cast, List, Optional, Tuple, Union
diff --git a/torch/optim/optimizer.py b/torch/optim/optimizer.py
index 129e377293f..49ed9049efc 100644
--- a/torch/optim/optimizer.py
+++ b/torch/optim/optimizer.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 """Base optimizer."""
 import functools
diff --git a/torch/optim/radam.py b/torch/optim/radam.py
index df0fbcef11a..23abbae80f6 100644
--- a/torch/optim/radam.py
+++ b/torch/optim/radam.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 r"""Implementation for the RAdam algorithm."""
 from typing import cast, List, Optional, Tuple, Union
diff --git a/torch/optim/rmsprop.py b/torch/optim/rmsprop.py
index d65b72088cf..853e4e231d3 100644
--- a/torch/optim/rmsprop.py
+++ b/torch/optim/rmsprop.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 r"""Implementation for the RMSprop algorithm."""
 from typing import List, Optional, Union
diff --git a/torch/optim/rprop.py b/torch/optim/rprop.py
index 100586fd6a2..d7b072c8992 100644
--- a/torch/optim/rprop.py
+++ b/torch/optim/rprop.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 r"""Implementation for the Resilient backpropagation."""
 from typing import List, Optional, Tuple, Union
diff --git a/torch/signal/windows/windows.py b/torch/signal/windows/windows.py
index f9f73b2dca0..51784b11194 100644
--- a/torch/signal/windows/windows.py
+++ b/torch/signal/windows/windows.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from typing import Optional, Iterable
 
diff --git a/torch/sparse/_triton_ops.py b/torch/sparse/_triton_ops.py
index 7585fc5e3e6..4ab313064ba 100644
--- a/torch/sparse/_triton_ops.py
+++ b/torch/sparse/_triton_ops.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import math
 import os
diff --git a/torch/testing/_internal/custom_op_db.py b/torch/testing/_internal/custom_op_db.py
index a2c0b439e7c..f15e8312aa5 100644
--- a/torch/testing/_internal/custom_op_db.py
+++ b/torch/testing/_internal/custom_op_db.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import torch
 import functools
diff --git a/torch/utils/_freeze.py b/torch/utils/_freeze.py
index f813ca28b81..d2d6cea7f22 100644
--- a/torch/utils/_freeze.py
+++ b/torch/utils/_freeze.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 """
 Freeze Python packages.
diff --git a/torch/utils/checkpoint.py b/torch/utils/checkpoint.py
index 1c0ff2f36fb..1a66e749de1 100644
--- a/torch/utils/checkpoint.py
+++ b/torch/utils/checkpoint.py
@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import contextlib
 import platform
diff --git a/torch/utils/flop_counter.py b/torch/utils/flop_counter.py
index a4f05c6c720..9b756f7635d 100644
--- a/torch/utils/flop_counter.py
+++ b/torch/utils/flop_counter.py
@@ -1,4 +1,5 @@
 # mypy: allow-untyped-defs
+# mypy: allow-untyped-decorators
 import torch
 from torch.utils._pytree import tree_map, tree_flatten, tree_unflatten
 from .module_tracker import ModuleTracker