mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
7ce69d5dbe
In almost all cases this is only included for writing the output formatter, which only uses `std::ostream` so including `<ostream>` is sufficient. The istream header is ~1000 lines so the difference is non-trivial. Pull Request resolved: https://github.com/pytorch/pytorch/pull/108150 Approved by: https://github.com/albanD, https://github.com/malfet ghstack dependencies: #108149
30 lines
515 B
C++
30 lines
515 B
C++
#include <torch/csrc/jit/runtime/print_handler.h>
|
|
|
|
#include <atomic>
|
|
#include <iostream>
|
|
#include <string>
|
|
|
|
namespace torch {
|
|
namespace jit {
|
|
|
|
namespace {
|
|
|
|
std::atomic<PrintHandler> print_handler(getDefaultPrintHandler());
|
|
|
|
} // namespace
|
|
|
|
PrintHandler getDefaultPrintHandler() {
|
|
return [](const std::string& s) { std::cout << s; };
|
|
}
|
|
|
|
PrintHandler getPrintHandler() {
|
|
return print_handler.load();
|
|
}
|
|
|
|
void setPrintHandler(PrintHandler ph) {
|
|
print_handler.store(ph);
|
|
}
|
|
|
|
} // namespace jit
|
|
} // namespace torch
|