pytorch

1
#include <torch/csrc/jit/mobile/model_tracer/CustomClassTracer.h>
2
#include <mutex>
3

4
namespace torch {
5
namespace jit {
6
namespace mobile {
7
CustomClassTracer::CustomClassTracer() {
8
  auto recorder_cb =
9
      [](const at::RecordFunction& fn) -> std::unique_ptr<at::ObserverContext> {
10
    std::string name = fn.name();
11
    getLoadedClasses().withLock(
12
        [&name](CustomClassTracer::custom_classes_type& custom_classes) {
13
          custom_classes.insert(name);
14
        });
15
    return nullptr;
16
  };
17

18
  handle_ = at::addGlobalCallback(at::RecordFunctionCallback(recorder_cb)
19
                                      .scopes({at::RecordScope::CUSTOM_CLASS}));
20
}
21

22
c10::Synchronized<CustomClassTracer::custom_classes_type>& CustomClassTracer::
23
    getLoadedClasses() {
24
  static c10::Synchronized<custom_classes_type> loaded_classes;
25
  return loaded_classes;
26
}
27

28
} // namespace mobile
29
} // namespace jit
30
} // namespace torch
31