mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 00:21:07 +01:00
bfeb45e46b
Follows #132604 Pull Request resolved: https://github.com/pytorch/pytorch/pull/132753 Approved by: https://github.com/Skylion007
29 lines
832 B
C++
29 lines
832 B
C++
#pragma once
|
|
|
|
#include <torch/csrc/jit/ir/ir.h>
|
|
#include <torch/csrc/onnx/onnx.h>
|
|
#include <torch/csrc/utils/pybind.h>
|
|
|
|
namespace torch::jit {
|
|
|
|
TORCH_API std::shared_ptr<Graph> ToONNX(
|
|
std::shared_ptr<Graph>& state,
|
|
::torch::onnx::OperatorExportTypes operator_export_type);
|
|
TORCH_API py::dict BlockToONNX(
|
|
Block* old_block,
|
|
Block* new_block,
|
|
::torch::onnx::OperatorExportTypes operator_export_type,
|
|
py::dict& env,
|
|
py::set& values_in_env,
|
|
bool is_sub_block = false);
|
|
TORCH_API void NodeToONNX(
|
|
Node* old_node,
|
|
Block* new_block,
|
|
::torch::onnx::OperatorExportTypes operator_export_type,
|
|
py::dict& env,
|
|
py::set& values_in_env);
|
|
TORCH_API void RemovePrintOps(std::shared_ptr<Graph>& graph);
|
|
TORCH_API void PreprocessCaffe2Ops(std::shared_ptr<Graph>& graph);
|
|
|
|
} // namespace torch::jit
|