torch.cuda.export_graph_data

torch.cuda.export_graph_data(path)

Return a post-instantiate hook that pickles CUDAGraph.get_graph_data() to path. Register it with CUDAGraph.register_post_instantiate_hook(): get_graph_data needs the graph instantiated (it remaps node ids to the exec graph id), and at post-instantiate time the template is still live, so this works for both keep_graph modes.

Return type:

Callable[[CUDAGraph], None]