1 | #ifndef THCP_STREAM_INC |
---|---|
2 | #define THCP_STREAM_INC |
3 | |
4 | #include <c10/cuda/CUDAStream.h> |
5 | #include <torch/csrc/Stream.h> |
6 | #include <torch/csrc/python_headers.h> |
7 | |
8 | struct THCPStream : THPStream { |
9 | at::cuda::CUDAStream cuda_stream; |
10 | }; |
11 | extern PyObject* THCPStreamClass; |
12 | |
13 | void THCPStream_init(PyObject* module); |
14 | |
15 | inline bool THCPStream_Check(PyObject* obj) { |
16 | return THCPStreamClass && PyObject_IsInstance(obj, THCPStreamClass); |
17 | } |
18 | |
19 | #endif // THCP_STREAM_INC |
20 |