torch.cuda.graphs.register_graph_capture_end_hook#

torch.cuda.graphs.register_graph_capture_end_hook(fn)[source]#

Register a hook run with each CUDA graph when its capture ends, while the captured cudaGraph_t is still live (see CUDAGraph.register_capture_end_hook()). Returns a RemovableHandle; call .remove() to unregister.