graph

class torch.cuda.graph(cuda_graph, pool=None, stream=None) [source]

Context-manager that captures CUDA work into a torch.cuda.CUDAGraph object for later replay.

See CUDA Graphs for a general introduction, detailed use, and constraints.

Parameters:

• cuda_graph (torch.cuda.CUDAGraph) – Graph object used for capture.
• pool (optional) – Opaque token (returned by a call to graph_pool_handle() or other_Graph_instance.pool()) hinting this graph’s capture may share memory from the specified pool. See Graph memory management.
• stream (torch.cuda.Stream, optional) – If supplied, will be set as the current stream in the context. If not supplied, graph sets its own internal side stream as the current stream in the context.

Note

For effective memory sharing, if you pass a pool used by a previous capture and the previous capture used an explicit stream argument, you should pass the same stream argument to this capture.

Warning

This API is in beta and may change in future releases.