Caffe2 - C++ API
A deep learning, cross platform ML framework
hooks_for_testing.cpp
1 #include <torch/csrc/jit/hooks_for_testing.h>
2 #include <torch/csrc/jit/script/module.h>
3 
4 namespace torch {
5 namespace jit {
6 
7 static std::function<void(std::shared_ptr<script::Module> module)>
8  emit_module_callback;
9 TORCH_API void didFinishEmitModule(std::shared_ptr<script::Module> module) {
10  if (emit_module_callback) {
11  emit_module_callback(std::move(module));
12  }
13 }
14 TORCH_API void setEmitModuleHook(
15  std::function<void(std::shared_ptr<script::Module> module)> cb) {
16  emit_module_callback = std::move(cb);
17 }
18 } // namespace jit
19 } // namespace torch
Definition: jit_type.h:17