| #pragma once |
| |
| #include <torch/csrc/jit/ir.h> |
| #include <torch/csrc/onnx/onnx.h> |
| |
| namespace torch { |
| namespace jit { |
| |
| TORCH_API std::shared_ptr<Graph> ToONNX( |
| std::shared_ptr<Graph>& state, |
| ::torch::onnx::OperatorExportTypes operator_export_type); |
| TORCH_API void BlockToONNX( |
| Block* old_block, |
| Block* new_block, |
| ::torch::onnx::OperatorExportTypes operator_export_type, |
| std::unordered_map<Value*, Value*> env); |
| TORCH_API void RemovePrintOps(std::shared_ptr<Graph>& graph); |
| TORCH_API void PreprocessCaffe2Ops(std::shared_ptr<Graph>& graph); |
| |
| } // namespace jit |
| } // namespace torch |
