[Unity] eliminate_dead_code on fx graph

python/tvm/relax/frontend/torch/dynamo.py

@@ -71,6 +71,8 @@ def to_tvm_tensor(torch_tensor):
             real_tensor = torch.randn(torch_tensor.shape, dtype=torch_tensor.dtype)
             return tvm.nd.array(real_tensor.numpy())
 
+        graph_module.graph.eliminate_dead_code()
+
         device = device_from_inputs(example_inputs)
 
         assert len(example_inputs)

tests/python/relax/test_frontend_dynamo.py

@@ -137,6 +137,21 @@ def forward(self, x):
         opt_model(inp).detach().numpy(), model(inp).detach().numpy(), rtol=1e-5, atol=1e-5
     )
 
+    def Func1(x, y):
+        z = torch.cat([x, y])
+        if z.size(0) > 5:
+            return z.mul(2)
+        else:
+            return z.add(2)
+
+    opt_func = torch.compile(Func1, backend=relax_dynamo(), dynamic=True)
+
+    for s in (2, 4):
+        x = torch.randn(s, 100)
+        y = torch.randn(s, 100)
+        with torch.no_grad():
+            tvm.testing.assert_allclose(opt_func(x, y), opt_func(x, y))
+
 
 def test_subgraph_capture():
     import torch
@@ -498,9 +513,7 @@ def main(
 
     class Select2(Module):
         def forward(self, input1):
-            result = input1[
-                torch.arange(1),
-            ]
+            result = input1[torch.arange(1)]
             return result
 
     verify_dynamo_model(