1 | #pragma once |
---|---|
2 | |
3 | #include <c10/core/SafePyObject.h> |
4 | #include <c10/macros/Macros.h> |
5 | #include <unordered_map> |
6 | |
7 | namespace at { |
8 | namespace impl { |
9 | |
10 | struct TORCH_API ThreadLocalPythonObjects { |
11 | static void set(const std::string& key, std::shared_ptr<SafePyObject> value); |
12 | static const std::shared_ptr<SafePyObject>& get(const std::string& key); |
13 | static bool contains(const std::string& key); |
14 | |
15 | static const ThreadLocalPythonObjects& get_state(); |
16 | static void set_state(ThreadLocalPythonObjects state); |
17 | |
18 | private: |
19 | std::unordered_map<std::string, std::shared_ptr<c10::SafePyObject>> obj_dict_; |
20 | }; |
21 | |
22 | } // namespace impl |
23 | } // namespace at |
24 |