1 files changed, 6 insertions, 3 deletions
diff --git a/src/error.rs b/src/error.rs
index 0114a86c..3f142960 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -1,4 +1,4 @@
-use crate::{DType, Device, Shape};
+use crate::{DType, DeviceLocation, Shape};
 
 /// Main library error type.
 #[derive(thiserror::Error, Debug)]
@@ -15,8 +15,8 @@ pub enum Error {
 
     #[error("device mismatch in {op}, lhs: {lhs:?}, rhs: {rhs:?}")]
     DeviceMismatchBinaryOp {
-        lhs: Device,
-        rhs: Device,
+        lhs: DeviceLocation,
+        rhs: DeviceLocation,
         op: &'static str,
     },
 
@@ -33,6 +33,9 @@ pub enum Error {
         got: usize,
         shape: Shape,
     },
+
+    #[error(transparent)]
+    Cuda(#[from] crate::CudaError),
 }
 
 pub type Result<T> = std::result::Result<T, Error>;