pytagi.cuda

Functions

is_available(→ bool)	Check if CUDA is available
is_nccl_available(→ bool)	Check if NCCL is available
get_device_count(→ int)	Get the number of CUDA devices
get_current_device(→ int)	Get the current CUDA device
set_device(→ bool)	Set the current CUDA device
is_device_available(→ bool)	Check if a specific CUDA device is available
get_device_properties(→ str)	Get the properties of a specific CUDA device
get_device_memory(→ Tuple[int, int])	Get the memory of a specific CUDA device

Module Contents

pytagi.cuda.is_available() → bool

Check if CUDA is available

Returns:

True if CUDA is available, False otherwise

Return type:

bool

pytagi.cuda.is_nccl_available() → bool

Check if NCCL is available

Returns:

True if NCCL is available, False otherwise

Return type:

bool

pytagi.cuda.get_device_count() → int

Get the number of CUDA devices

Returns:

Number of CUDA devices

Return type:

int

pytagi.cuda.get_current_device() → int

Get the current CUDA device

Returns:

Current CUDA device

Return type:

int

pytagi.cuda.set_device(device_index: int) → bool

Set the current CUDA device

Parameters:

device_index – Device index to set

pytagi.cuda.is_device_available(device_index: int) → bool

Check if a specific CUDA device is available

Parameters:

device_index – Device index to check

pytagi.cuda.get_device_properties(device_index: int) → str

Get the properties of a specific CUDA device

Parameters:

device_index – Device index to get properties of

pytagi.cuda.get_device_memory(device_index: int) → Tuple[int, int]

Get the memory of a specific CUDA device

Parameters:

device_index – Device index to get memory of