kyro 0.1.1

A high-performance ML inference engine
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

use candle_core::{Device, Result};

pub fn get_device() -> Result<Device> {
    if candle_core::utils::cuda_is_available() {
        Ok(Device::new_cuda(0)?)
    } else if candle_core::utils::metal_is_available() {
        Ok(Device::new_metal(0)?)
    } else {
        #[cfg(all(target_os = "macos", target_arch = "aarch64"))]
        {
            Ok(Device::new_metal(0)?)
        }
        #[cfg(not(all(target_os = "macos", target_arch = "aarch64")))]
        {
            Ok(Device::Cpu)
        }
    }
}