| 39 | |
| 40 | |
| 41 | class M(torch.nn.Module): |
| 42 | def __init__(self): |
| 43 | super().__init__() |
| 44 | |
| 45 | def forward(self, x, pred1, pred2, y): |
| 46 | def true_fn(x, pred2): |
| 47 | def true_nested(y): |
| 48 | y = y + y |
| 49 | y = torch.mm(y, y) |
| 50 | return y |
| 51 | |
| 52 | def false_nested(y): |
| 53 | return torch.mm(y, y) |
| 54 | |
| 55 | z = control_flow.cond(pred2, true_nested, false_nested, [x]) |
| 56 | return x + z |
| 57 | |
| 58 | def false_fn(x, _pred2): |
| 59 | return torch.mm(x, x) |
| 60 | |
| 61 | x = x.cos() |
| 62 | x = x + y |
| 63 | y = control_flow.cond(pred1, true_fn, false_fn, [x, pred2]) |
| 64 | return y.sin() |
| 65 | |
| 66 | def get_example_inputs(self): |
| 67 | return ( |
| 68 | torch.ones(2, 2), |
| 69 | torch.tensor([False]), |
| 70 | torch.Tensor([False]), |
| 71 | torch.ones(2, 2), |
| 72 | ) |
| 73 | |
| 74 | |
| 75 | @final |