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