torch.cuda.export_graph_data

torch.cuda.export_graph_data(path) [source]

Возвращает хук, вызываемый после инстанцирования, который сериализует с помощью pickle CUDAGraph.get_graph_data() в path. Зарегистрируйте его с помощью CUDAGraph.register_post_instantiate_hook(): для get_graph_data необходимо инстанцировать граф (он сопоставляет идентификаторы узлов с идентификатором графа exec), а на этапе после инстанцирования шаблон всё ещё активен, поэтому это работает в обоих режимах keep_graph.

Тип возвращаемого значения:

Callable[[CUDAGraph], None]