1 #pragma once 2 #include <torch/csrc/Export.h> 3 #include <memory> 4 #include <ostream> 5 #include <string> 6 7 namespace torch::jit::onnx { 8 9 TORCH_API bool is_log_enabled(); 10 11 TORCH_API void set_log_enabled(bool enabled); 12 13 TORCH_API void set_log_output_stream(std::shared_ptr<std::ostream> out_stream); 14 15 TORCH_API std::ostream& _get_log_output_stream(); 16 17 #define ONNX_LOG(...) \ 18 if (::torch::jit::onnx::is_log_enabled()) { \ 19 ::torch::jit::onnx::_get_log_output_stream() \ 20 << ::c10::str(__VA_ARGS__) << std::endl; \ 21 } 22 23 } // namespace torch::jit::onnx 24