1 | #pragma once |
2 | |
3 | #include <c10/core/TensorOptions.h> |
4 | |
5 | // cuda_lazy_init() is always compiled, even for CPU-only builds. |
6 | // Thus, it does not live in the cuda/ folder. |
7 | |
8 | namespace torch { |
9 | namespace utils { |
10 | |
11 | // The INVARIANT is that this function MUST be called before you attempt |
12 | // to get a CUDA Type object from ATen, in any way. Here are some common |
13 | // ways that a Type object may be retrieved: |
14 | // |
15 | // - You call getNonVariableType or getNonVariableTypeOpt |
16 | // - You call toBackend() on a Type |
17 | // |
18 | // It's important to do this correctly, because if you forget to add it |
19 | // you'll get an oblique error message about "Cannot initialize CUDA without |
20 | // ATen_cuda library" if you try to use CUDA functionality from a CPU-only |
21 | // build, which is not good UX. |
22 | // |
23 | void cuda_lazy_init(); |
24 | void set_requires_cuda_init(bool value); |
25 | |
26 | static void maybe_initialize_cuda(const at::TensorOptions& options) { |
27 | if (options.device().is_cuda()) { |
28 | torch::utils::cuda_lazy_init(); |
29 | } |
30 | } |
31 | |
32 | } // namespace utils |
33 | } // namespace torch |
34 | |