torch.fx.passes.backends.cudagraphs.partition_cudagraphs#

torch.fx.passes.backends.cudagraphs.partition_cudagraphs(gm, inputs)[source]#

Partition an FX graph into sub-GraphModules that can be validly run under CUDA graphs. For a subgraph to be runnable under CUDA, all of the operations must involve CUDA tensors only/