add a test for inline tracing

Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/23535 Test Plan: Imported from OSS Differential Revision: D16554694 Pulled By: suo fbshipit-source-id: 0fae4458f18c06ffbd484905ad7836dce9ce69cc

add a test for inline tracing
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/23535 Test Plan: Imported from OSS Differential Revision: D16554694 Pulled By: suo fbshipit-source-id: 0fae4458f18c06ffbd484905ad7836dce9ce69cc
265b498d · Michael Suo · Facebook Github Bot · 52b95fd4 · 265b498d
--- a/test/test_jit.py
+++ b/test/test_jit.py
@@ -329,6 +329,22 @@ class TestJit(JitTestCase):
        t2 = torch.rand(3, 4)
        self.assertEqual(r(t2), m.x + t2)

+    def test_trace_nested_fn(self):
+        class TracedInlineDecision(torch.nn.Module):
+            def forward(self, x, flag):
+                @torch.jit.script
+                def make_decision(flag, x):
+                    if flag:
+                        return x
+                    else:
+                        return torch.zeros_like(x)
+                x = torch.neg(x) # do something random
+                return make_decision(flag, x)
+
+
+        decision = TracedInlineDecision()
+        torch.jit.trace(decision, (torch.rand(3, 4), torch.tensor([True], dtype=torch.bool)), check_trace=True)
+
    def test_restore_device(self):
        # main purpose is checking map_location works
        m = torch.jit.ScriptModule()