adapt to new metal device enum

mimiquate · Nov 13, 2023 · 3f00264 · 3f00264
1 parent 2a905cd
commit 3f00264
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/native/candlex/src/lib.rs b/native/candlex/src/lib.rs
@@ -1,7 +1,8 @@
 mod atoms {
     rustler::atoms! {
         cpu,
-        cuda
+        cuda,
+        metal
     }
 }
 

diff --git a/native/candlex/src/tensors.rs b/native/candlex/src/tensors.rs
@@ -26,6 +26,7 @@ impl ExTensor {
         let dev_string = match tensor.device() {
             Device::Cpu => atoms::cpu(),
             Device::Cuda(_) => atoms::cuda(),
+            Device::Metal(_) => atoms::metal(),
         };
 
         Self {