mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
ad8aef0f98
Mostly in torch/csrc/jit/runtime and in `ATen/cuda/` Pull Request resolved: https://github.com/pytorch/pytorch/pull/110314 Approved by: https://github.com/seemethere
28 lines
490 B
C++
28 lines
490 B
C++
#include <torch/csrc/jit/runtime/print_handler.h>
|
|
|
|
#include <atomic>
|
|
#include <iostream>
|
|
#include <string>
|
|
|
|
namespace torch::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 torch::jit
|