Update on "[NJT] Inline through torch.nested.nested_tensor_from_jagged instead of graph break"

soulitzer · soulitzer · commit 53d8ecf21f92 · 2024-04-19T14:15:27.000-04:00
cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang

[ghstack-poisoned]
diff --git a/test/profiler/test_profiler.py b/test/profiler/test_profiler.py
@@ -650,6 +650,7 @@ def test_execution_trace_no_capture(self):
                 found_root_node = True
         assert found_root_node
 
+    @skipIfTorchDynamo("https://github.com/pytorch/pytorch/issues/124500")
     def test_execution_trace_nested_tensor(self):
         fp = tempfile.NamedTemporaryFile("w+t", suffix=".et.json", delete=False)
         fp.close()
@@ -664,9 +665,6 @@ def fn(nt):
             for i in range(3):
                 values = torch.rand((8 + i, 4 + i))
                 offsets = torch.tensor([0, 2, 4, 6, 8 + i])
-                # Need to explicitly graph break due to:
-                # https://github.com/pytorch/pytorch/issues/124500
-                torch._dynamo.graph_break()
                 nt = torch.nested.nested_tensor_from_jagged(values, offsets)
                 fn(nt)
 
