| 9254 | """Nested cond: a cond inside one of the branches.""" |
| 9255 | |
| 9256 | class CondNestedModel(Module): |
| 9257 | def forward(self, x): |
| 9258 | def true_fn(x): |
| 9259 | def inner_true(x): |
| 9260 | return x * 2.0 |
| 9261 | |
| 9262 | def inner_false(x): |
| 9263 | return x * 3.0 |
| 9264 | |
| 9265 | return torch.cond(x.sum() > 1, inner_true, inner_false, (x,)) |
| 9266 | |
| 9267 | def false_fn(x): |
| 9268 | return x - 1.0 |
| 9269 | |
| 9270 | return torch.cond(x.sum() > 0, true_fn, false_fn, (x,)) |
| 9271 | |
| 9272 | example_args = (torch.randn(3, 4),) |
| 9273 | exported_program = export(CondNestedModel(), args=example_args) |
no outgoing calls
searching dependent graphs…