ThreadLocalPythonObjects.h 861 B

1234567891011121314151617181920212223242526
  1. #if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
  2. #pragma once
  3. #include <c10/core/SafePyObject.h>
  4. #include <c10/macros/Macros.h>
  5. #include <unordered_map>
  6. namespace at::impl {
  7. struct TORCH_API ThreadLocalPythonObjects {
  8. static void set(const std::string& key, std::shared_ptr<SafePyObject> value);
  9. static const std::shared_ptr<SafePyObject>& get(const std::string& key);
  10. static bool contains(const std::string& key);
  11. static const ThreadLocalPythonObjects& get_state();
  12. static void set_state(ThreadLocalPythonObjects state);
  13. private:
  14. std::unordered_map<std::string, std::shared_ptr<c10::SafePyObject>> obj_dict_;
  15. };
  16. } // namespace at::impl
  17. #else
  18. #error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
  19. #endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)