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