#include "torch/csrc/python_headers.h" #include #include #include "THCP.h" #include "override_macros.h" #define THC_GENERIC_FILE "torch/csrc/generic/utils.cpp" #include #ifdef USE_CUDA std::vector THPUtils_PySequence_to_THCStreamList(PyObject *obj) { if (!PySequence_Check(obj)) { throw std::runtime_error("Expected a sequence in THPUtils_PySequence_to_THCStreamList"); } THPObjectPtr seq = THPObjectPtr(PySequence_Fast(obj, NULL)); if (seq.get() == NULL) { throw std::runtime_error("expected PySequence, but got " + std::string(THPUtils_typename(obj))); } std::vector streams; Py_ssize_t length = PySequence_Fast_GET_SIZE(seq.get()); for (Py_ssize_t i = 0; i < length; i++) { PyObject *stream = PySequence_Fast_GET_ITEM(seq.get(), i); if (PyObject_IsInstance(stream, THCPStreamClass)) { streams.push_back( ((THCPStream *)stream)->cdata); } else if (stream == Py_None) { streams.push_back(NULL); } else { std::runtime_error("Unknown data type found in stream list. Need THCStream or None"); } } return streams; } #endif