| 55 | |
| 56 | def test_exprimental_simplify_subgraph(): |
| 57 | class WithSubGraph(torch.nn.Module): |
| 58 | def __init__(self): |
| 59 | super(WithSubGraph, self).__init__() |
| 60 | |
| 61 | def forward(self, x): |
| 62 | if x.sum() > 1.0: |
| 63 | # NOTE: even onnxsim cannot simplify it, |
| 64 | # a canonical pass in onnx-optimizer is needed for it. |
| 65 | # so this test only tests that include_subgraph doesn't |
| 66 | # result in invalid model in this case |
| 67 | return 3 + x + 3 |
| 68 | else: |
| 69 | return x + 4 |
| 70 | |
| 71 | net = torch.jit.script(WithSubGraph()) |
| 72 | dummy_input = torch.randn(2) |
no outgoing calls