ThreadLocalPythonObjects.h 607 B

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