mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
5741de883a
Summary: The record_stream method was hard coded for CUDA device. Define the record_stream in the native_functions.yaml to enable the dynamic dispatch to different end device. Fixes https://github.com/pytorch/pytorch/issues/36556 Pull Request resolved: https://github.com/pytorch/pytorch/pull/44301 Reviewed By: glaringlee Differential Revision: D23763954 Pulled By: ezyang fbshipit-source-id: e6d24f5e7892b56101fa858a6cad2abc5cdc4293
19 lines
374 B
C
19 lines
374 B
C
#ifndef THP_STREAM_INC
|
|
#define THP_STREAM_INC
|
|
|
|
#include <torch/csrc/python_headers.h>
|
|
|
|
struct THPStream {
|
|
PyObject_HEAD
|
|
uint64_t cdata;
|
|
};
|
|
extern PyTypeObject *THPStreamClass;
|
|
|
|
void THPStream_init(PyObject *module);
|
|
|
|
inline bool THPStream_Check(PyObject* obj) {
|
|
return THPStreamClass && PyObject_IsInstance(obj, (PyObject*)THPStreamClass);
|
|
}
|
|
|
|
#endif // THP_STREAM_INC
|