1 #pragma once 2 #include <torch/csrc/Export.h> 3 #include <torch/csrc/jit/api/compilation_unit.h> 4 #include <functional> 5 #include <memory> 6 7 namespace torch { 8 namespace jit { 9 struct Module; 10 11 using ModuleHook = std::function<void(Module module)>; 12 using FunctionHook = std::function<void(StrongFunctionPtr function)>; 13 14 TORCH_API void didFinishEmitModule(Module module); 15 TORCH_API void didFinishEmitFunction(StrongFunctionPtr defined); 16 TORCH_API void setEmitHooks(ModuleHook for_module, FunctionHook for_fn); 17 18 TORCH_API std::pair<ModuleHook, FunctionHook> getEmitHooks(); 19 20 } // namespace jit 21 } // namespace torch 22