1 #ifndef THP_STREAM_INC 2 #define THP_STREAM_INC 3 4 #include <c10/core/Stream.h> 5 #include <c10/macros/Export.h> 6 #include <torch/csrc/python_headers.h> 7 8 struct THPStream { 9 PyObject_HEAD int64_t stream_id; 10 int64_t device_type; 11 int64_t device_index; 12 }; 13 extern TORCH_API PyTypeObject* THPStreamClass; 14 15 void THPStream_init(PyObject* module); 16 THPStream_Check(PyObject * obj)17inline bool THPStream_Check(PyObject* obj) { 18 return THPStreamClass && PyObject_IsInstance(obj, (PyObject*)THPStreamClass); 19 } 20 21 PyObject* THPStream_Wrap(const c10::Stream& stream); 22 23 #endif // THP_STREAM_INC 24