module1 = torch.cuda.make_graphed_callables(module1, (x,))
module2 = torch.cuda.make_graphed_callables(module2, (h,))
module3 = torch.cuda.make_graphed_callables(module3, (h,))
real_inputs = [torch.rand_like(x) for _ in range(10)]
real_targets = [torch.randn(N, D_out, device="cuda") for _ in range(10)]
for data, target in zip(real_inputs, real_targets):
    optimizer.zero_grad(set_to_none=True)
tmp = module1(data)  # forward ops run as a graph

if tmp.sum().item() > 0:
    tmp = module2(tmp)  # forward ops run as a graph
else:
    tmp = module3(tmp)  # forward ops run as a graph

loss = loss_fn(tmp, target)
# module2's or module3's (whichever was chosen) backward ops,
# as well as module1's backward ops, run as graphs
loss.backward()
optimizer.step()

```
Example use cases
MLPerf v1.0 training workloads