xref: /aosp_15_r20/external/pytorch/torch/csrc/jit/runtime/print_handler.cpp (revision da0073e96a02ea20f0ac840b70461e3646d07c45) 
1 #include <torch/csrc/jit/runtime/print_handler.h>
2 
3 #include <atomic>
4 #include <iostream>
5 #include <string>
6 
7 namespace torch::jit {
8 
9 namespace {
10 
11 std::atomic<PrintHandler> print_handler(getDefaultPrintHandler());
12 
13 } // namespace
14 
getDefaultPrintHandler()15 PrintHandler getDefaultPrintHandler() {
16   return [](const std::string& s) { std::cout << s; };
17 }
18 
getPrintHandler()19 PrintHandler getPrintHandler() {
20   return print_handler.load();
21 }
22 
setPrintHandler(PrintHandler ph)23 void setPrintHandler(PrintHandler ph) {
24   print_handler.store(ph);
25 }
26 
27 } // namespace torch::jit
28