Enum tch::Cuda

pub enum Cuda {}

Cuda related helper functions.

Implementations

impl Cuda

pub fn device_count() -> i64

Returns the number of CUDA devices available.

pub fn is_available() -> bool

Returns true if at least one CUDA device is available.

pub fn cudnn_is_available() -> bool

Returns true if CUDA is available, and CuDNN is available.

pub fn manual_seed(seed: u64)

Sets the seed for the current GPU.

Arguments

• seed - An unsigned 64bit int to be used as seed.

pub fn manual_seed_all(seed: u64)

Sets the seed for all available GPUs.

Arguments

• seed - An unsigned 64bit int to be used as seed.

pub fn synchronize(device_index: i64)

Waits for all kernels in all streams on a CUDA device to complete.

Arguments

• device_index - A signed 64bit int to indice which device to wait for.

pub fn user_enabled_cudnn() -> bool

Returns true if cudnn is enabled by the user.

This does not indicate whether cudnn is actually usable.

pub fn set_user_enabled_cudnn(b: bool)

Enable or disable cudnn.

pub fn cudnn_set_benchmark(b: bool)

Sets cudnn benchmark mode.

When set cudnn will try to optimize the generators durning the first network runs and then use the optimized architecture in the following runs. This can result in significant performance improvements.