huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-16 10:38:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

Metal part 1 - Scaffolding for metal. (#1308)

Browse Source 
* Metal part 1 - Scaffolding for metal.

* Remove tracing.
This commit is contained in:
Nicolas Patry
2023-11-10 08:35:48 +01:00
committed by GitHub
parent 18d30005c5
commit 26c4e5bf1d
13 changed files with 473 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
candle-examples/src/lib.rs
View File

@ -2,17 +2,28 @@ pub mod coco_classes;
pub mod imagenet;
pub mod token_output_stream;
use candle::utils::{cuda_is_available, metal_is_available};
use candle::{Device, Result, Tensor};
pub fn device(cpu: bool) -> Result<Device> {
if cpu {
Ok(Device::Cpu)
} else if cuda_is_available() {
Ok(Device::new_cuda(0)?)
} else if metal_is_available() {
Ok(Device::new_metal(0)?)
} else {
let device = Device::cuda_if_available(0)?;
if !device.is_cuda() {
#[cfg(all(target_os = "macos", target_arch = "aarch64"))]
{
println!(
"Running on CPU, to run on GPU(metal), build this example with `--features metal`"
);
}
#[cfg(not(all(target_os = "macos", target_arch = "aarch64")))]
{
println!("Running on CPU, to run on GPU, build this example with `--features cuda`");
}
Ok(device)
Ok(Device::Cpu)
}
}