1#include <torch/csrc/utils/cuda_lazy_init.h>
2
3#include <torch/csrc/Exceptions.h>
4#include <torch/csrc/python_headers.h>
5#include <torch/csrc/utils/object_ptr.h>
6
7namespace torch {
8namespace utils {
9namespace {
10
11bool is_initialized = false;
12
13}
14
15void cuda_lazy_init() {
16 pybind11::gil_scoped_acquire g;
17 // Protected by the GIL. We don't use call_once because under ASAN it
18 // has a buggy implementation that deadlocks if an instance throws an
19 // exception. In any case, call_once isn't necessary, because we
20 // have taken a lock.
21 if (is_initialized) {
22 return;
23 }
24
25 auto module = THPObjectPtr(PyImport_ImportModule("torch.cuda"));
26 if (!module) {
27 throw python_error();
28 }
29
30 auto res = THPObjectPtr(PyObject_CallMethod(module.get(), "_lazy_init", ""));
31 if (!res) {
32 throw python_error();
33 }
34
35 is_initialized = true;
36}
37
38void set_requires_cuda_init(bool value) {
39 is_initialized = !value;
40}
41
42} // namespace utils
43} // namespace torch
44