mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
20 lines
374 B
C
20 lines
374 B
C
#pragma once
|
|
|
|
#include "torch/csrc/python_headers.h"
|
|
#include "torch/csrc/utils/device.h"
|
|
|
|
struct THPDevice {
|
|
PyObject_HEAD
|
|
torch::Device device;
|
|
};
|
|
|
|
extern PyTypeObject THPDeviceType;
|
|
|
|
inline bool THPDevice_Check(PyObject *obj) {
|
|
return Py_TYPE(obj) == &THPDeviceType;
|
|
}
|
|
|
|
PyObject * THPDevice_New(const torch::Device& device);
|
|
|
|
void THPDevice_init(PyObject *module);
|