mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
c2a18a6702
Summary: This makes it so we can see the output of prim::Print in environments like iPython notebooks which override sys.stdout Pull Request resolved: https://github.com/pytorch/pytorch/pull/21625 Differential Revision: D15756793 Pulled By: jamesr66a fbshipit-source-id: 7d9a14b2e229ed358e784318e9d862677db2c461
23 lines
388 B
C++
23 lines
388 B
C++
#include <torch/csrc/jit/print_handler.h>
|
|
|
|
#include <iostream>
|
|
#include <string>
|
|
|
|
namespace torch {
|
|
namespace jit {
|
|
|
|
std::atomic<PrintHandler> print_handler([](const std::string& str) {
|
|
std::cout << str;
|
|
});
|
|
|
|
PrintHandler getPrintHandler() {
|
|
return print_handler.load();
|
|
}
|
|
|
|
void setPrintHandler(PrintHandler ph) {
|
|
print_handler.store(ph);
|
|
}
|
|
|
|
} // namespace jit
|
|
} // namespace torch
|