Enum Device

pub enum Device {
    Cpu,
    Cuda(CudaDevice),
    Metal(MetalDevice),
}

Variants

Cpu

Cuda(CudaDevice)

Metal(MetalDevice)

Implementations

impl Device

pub fn new_cuda(ordinal: usize) -> Result<Device, Error>

pub fn as_cuda_device(&self) -> Result<&CudaDevice, Error>

pub fn as_metal_device(&self) -> Result<&MetalDevice, Error>

pub fn new_cuda_with_stream(ordinal: usize) -> Result<Device, Error>

pub fn new_metal(ordinal: usize) -> Result<Device, Error>

pub fn set_seed(&self, seed: u64) -> Result<(), Error>

pub fn get_current_seed(&self) -> Result<u64, Error>

Get the current seed for the device RNG.

pub fn same_device(&self, rhs: &Device) -> bool

pub fn location(&self) -> DeviceLocation

pub fn is_cpu(&self) -> bool

pub fn is_cuda(&self) -> bool

pub fn is_metal(&self) -> bool

pub fn supports_bf16(&self) -> bool

pub fn bf16_default_to_f32(&self) -> DType

Return BF16 for devices that support it, otherwise default to F32.

pub fn cuda_if_available(ordinal: usize) -> Result<Device, Error>

pub fn synchronize(&self) -> Result<(), Error>

Trait Implementations

impl Clone for Device

fn clone(&self) -> Device

Returns a copy of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for Device

fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), Error>

Formats the value using the given formatter.