1 #pragma once 2 3 #include <torch/csrc/jit/ir/ir.h> 4 #include <torch/csrc/onnx/onnx.h> 5 #include <torch/csrc/utils/pybind.h> 6 7 namespace torch::jit { 8 9 TORCH_API std::shared_ptr<Graph> ToONNX( 10 std::shared_ptr<Graph>& state, 11 ::torch::onnx::OperatorExportTypes operator_export_type); 12 TORCH_API py::dict BlockToONNX( 13 Block* old_block, 14 Block* new_block, 15 ::torch::onnx::OperatorExportTypes operator_export_type, 16 py::dict& env, 17 py::set& values_in_env, 18 bool is_sub_block = false); 19 TORCH_API void NodeToONNX( 20 Node* old_node, 21 Block* new_block, 22 ::torch::onnx::OperatorExportTypes operator_export_type, 23 py::dict& env, 24 py::set& values_in_env); 25 TORCH_API void RemovePrintOps(std::shared_ptr<Graph>& graph); 26 TORCH_API void PreprocessCaffe2Ops(std::shared_ptr<Graph>& graph); 27 28 } // namespace torch::jit 29