1 | #pragma once |
---|---|
2 | |
3 | #include <torch/csrc/Export.h> |
4 | #include <torch/csrc/python_headers.h> |
5 | |
6 | #include <ATen/Device.h> |
7 | |
8 | struct TORCH_API THPDevice { |
9 | PyObject_HEAD at::Device device; |
10 | }; |
11 | |
12 | TORCH_API extern PyTypeObject THPDeviceType; |
13 | |
14 | inline bool THPDevice_Check(PyObject* obj) { |
15 | return Py_TYPE(obj) == &THPDeviceType; |
16 | } |
17 | |
18 | TORCH_API PyObject* THPDevice_New(const at::Device& device); |
19 | |
20 | TORCH_API void THPDevice_init(PyObject* module); |
21 |