Merge pull request #3 from LaurentMazare/cuda

Add Cuda support.

1 files changed, 10 insertions, 1 deletions

diff --git a/src/lib.rs b/src/lib.rs
index 175d36ad..3bae1a7e 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,6 +1,9 @@
 mod cpu_backend;
+#[cfg(feature = "cuda")]
+mod cuda_backend;
 mod device;
 mod dtype;
+mod dummy_cuda_backend;
 mod error;
 mod op;
 mod shape;
@@ -9,10 +12,16 @@ mod strided_index;
 mod tensor;
 
 pub use cpu_backend::CpuStorage;
-pub use device::Device;
+pub use device::{Device, DeviceLocation};
 pub use dtype::{DType, WithDType};
 pub use error::{Error, Result};
 pub use shape::Shape;
 pub use storage::Storage;
 use strided_index::StridedIndex;
 pub use tensor::{Tensor, TensorId};
+
+#[cfg(feature = "cuda")]
+pub use cuda_backend::{CudaDevice, CudaError, CudaStorage};
+
+#[cfg(not(feature = "cuda"))]
+pub use dummy_cuda_backend::{CudaDevice, CudaError, CudaStorage};