Make TORCH_COMPILE_DEBUG=1 work again

pytorch · Nov 3, 2023 · 2016fc9 · 2016fc9
1 parent 132cb57
commit 2016fc9
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/test/inductor/test_torchinductor.py b/test/inductor/test_torchinductor.py
@@ -3111,6 +3111,7 @@ def forward(self, l_input_: torch.Tensor):
         o2 = torch.compile(mod)(inp)
         self.assertEqual(o1, o2)
 
+    @patch.object(config.trace, "enabled", True)
     def test_layer_norm(self):
         m = torch.nn.Sequential(
             torch.nn.LayerNorm(32),

diff --git a/torch/_inductor/scheduler.py b/torch/_inductor/scheduler.py
@@ -779,7 +779,7 @@ def __init__(self, scheduler: "Scheduler", snodes: List[SchedulerNode]):
         # NB: No need to call super().__init__() because we don't need to re-use any of its logic.
         self.snodes = snodes
         self.scheduler = scheduler
-        self.node: ir.Buffer
+        self.node: Optional[ir.Buffer] = None
         self.users: List[NodeUser] = []
         self.inverse_users = []
         self.node_users = []
@@ -1017,7 +1017,7 @@ def __init__(
         else:
             self.scheduler = scheduler
             self.snodes = nodes
-            self.node: ir.Buffer
+            self.node: Optional[ir.Buffer] = None
             self.users: List[NodeUser] = []
 
             self.set_read_writes(