torch.cuda.make_graphed_callables
torch.cuda.make_graphed_callables(callables, sample_args, num_warmup_iters=3, allow_unused_input=False)
Accepts callables (functions or "nn.Module"s) and returns graphed
   versions.
Each graphed callable's forward pass runs its source callable's
   forward CUDA work as a CUDA graph inside a single autograd node.
The graphed callable's forward pass also appends a backward node to
   the autograd graph. During backward, this node runs the callable's
   backward work as a CUDA graph.
Therefore, each graphed callable should be a drop-in replacement
   for its source callable in an autograd-enabled training loop.
See Partial-network capture for detailed use and constraints.
If you pass a tuple of several callables, their captures will use
   the same memory pool. See Graph memory management for when this is
   appropriate.
Parameters:
      * callables (torch.nn.Module or Python function*, or