Rate this Page
★★★★★

graph#

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

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

SeeCUDA 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 tograph_pool_handle() orother_Graph_instance.pool()) hinting this graph’s capturemay share memory from the specified pool. SeeGraph 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.

  • capture_error_mode (str,optional) – specifies the cudaStreamCaptureMode for the graph capture stream.Can be “global”, “thread_local” or “relaxed”. During cuda graph capture, some actions, such as cudaMalloc,may be unsafe. “global” will error on actions in other threads, “thread_local” will only error foractions in the current thread, and “relaxed” will not error on actions. Do NOT change this settingunless you’re familiar withcudaStreamCaptureMode

Note

For effective memory sharing, if you pass apool used by a previous capture and the previous captureused an explicitstream argument, you should pass the samestream argument to this capture.

Warning

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