1 | #pragma once |
2 | |
3 | #include <torch/csrc/Export.h> |
4 | |
5 | #include <cstddef> |
6 | #include <cstdint> |
7 | |
8 | namespace torch { |
9 | namespace cuda { |
10 | |
11 | /// Returns the number of CUDA devices available. |
12 | size_t TORCH_API device_count(); |
13 | |
14 | /// Returns true if at least one CUDA device is available. |
15 | bool TORCH_API is_available(); |
16 | |
17 | /// Returns true if CUDA is available, and CuDNN is available. |
18 | bool TORCH_API cudnn_is_available(); |
19 | |
20 | /// Sets the seed for the current GPU. |
21 | void TORCH_API manual_seed(uint64_t seed); |
22 | |
23 | /// Sets the seed for all available GPUs. |
24 | void TORCH_API manual_seed_all(uint64_t seed); |
25 | |
26 | /// Waits for all kernels in all streams on a CUDA device to complete. |
27 | void TORCH_API synchronize(int64_t device_index = -1); |
28 | |
29 | } // namespace cuda |
30 | } // namespace torch |
31 | |