Add some 'cuda-if-available' helper function. (#172)

1 files changed, 12 insertions, 0 deletions

diff --git a/candle-examples/src/lib.rs b/candle-examples/src/lib.rs
index 8b137891..285aee04 100644
--- a/candle-examples/src/lib.rs
+++ b/candle-examples/src/lib.rs
@@ -1 +1,13 @@
+use candle::{Device, Result};
 
+pub fn device(cpu: bool) -> Result<Device> {
+    if cpu {
+        Ok(Device::Cpu)
+    } else {
+        let device = Device::cuda_if_available(0)?;
+        if !device.is_cuda() {
+            println!("Running on CPU, to run on GPU, build this example with `--features cuda`");
+        }
+        Ok(device)
+    }
+}