diff --git a/Cargo.toml b/Cargo.toml index fb504d7..ea81937 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -36,6 +36,7 @@ rayon = "1.10" [features] flash-attn=["candle-flash-attn"] cuda=["candle-nn/cuda", "candle-core/cuda", "candle-transformers/cuda"] +metal=["candle-nn/metal", "candle-core/metal", "candle-transformers/metal"] ffmpeg=["ffmpeg-next"] [lints.clippy] diff --git a/src/utils/mod.rs b/src/utils/mod.rs index d14996e..32995f0 100644 --- a/src/utils/mod.rs +++ b/src/utils/mod.rs @@ -27,7 +27,11 @@ pub fn get_device(device: Option<&Device>) -> Device { { Device::new_cuda(0).unwrap_or(Device::Cpu) } - #[cfg(not(feature = "cuda"))] + #[cfg(all(not(feature = "cuda"), feature = "metal"))] + { + Device::new_metal(0).unwrap_or(Device::Cpu) + } + #[cfg(all(not(feature = "cuda"), not(feature = "metal")))] { Device::Cpu }