hooks_for_testing.cpp

#include "torch/csrc/jit/hooks_for_testing.h"
#include "torch/csrc/jit/script/module.h"

namespace torch {
namespace jit {

static std::function<void(std::shared_ptr<script::Module> module)> emit_module_callback;
TORCH_API void didFinishEmitModule(std::shared_ptr<script::Module> module) {
  if(emit_module_callback)
    emit_module_callback(module);

}
TORCH_API void setEmitModuleHook(std::function<void(std::shared_ptr<script::Module> module)> cb) {
  emit_module_callback = cb;
}
} // namespace jit
} // namespace torch