tensorplay.cuda.graphs.graph

tensorplay.cuda.graphs.graph(cuda_graph, pool=None, stream=None, capture_error_mode='global')[source]

Context-manager that captures CUDA work into a tensorplay.cuda.CUDAGraph.

Parameters:

• cuda_graph (CUDAGraph) – the graph object to capture into.

• pool – None gives the graph its own private memory pool; an id from graph_pool_handle() (or another graph / its pool_id property) shares that pool so allocations from both captures can recycle each other’s space.

• stream (Stream, optional) – custom capture stream; defaults to the runtime’s dedicated per-device side stream (the legacy default stream cannot capture).

• capture_error_mode (str, optional) – see CUDAGraph.capture_begin().