xref: /aosp_15_r20/external/pytorch/torch/csrc/Device.h (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
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 // NOLINTNEXTLINE(cppcoreguidelines-pro-type-member-init)
9 struct TORCH_API THPDevice {
10   PyObject_HEAD at::Device device;
11 };
12 
13 TORCH_API extern PyTypeObject THPDeviceType;
14 
THPDevice_Check(PyObject * obj)15 inline bool THPDevice_Check(PyObject* obj) {
16   return Py_TYPE(obj) == &THPDeviceType;
17 }
18 
19 TORCH_API PyObject* THPDevice_New(const at::Device& device);
20 
21 TORCH_API void THPDevice_init(PyObject* module);
22