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

Move the data between the host and the device.

Browse Source
This commit is contained in:
laurent
2023-06-21 19:43:25 +01:00
parent c654ecdb16
commit 71735c7a02
5 changed files with 100 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
src/storage.rs
View File

@ -1,9 +1,9 @@
use crate::{CpuStorage, DType, Device, Error, Result, Shape};
use crate::{CpuStorage, CudaStorage, DType, Device, Error, Result, Shape};
#[derive(Debug, Clone)]
pub enum Storage {
Cpu(CpuStorage),
Cuda(cudarc::driver::CudaSlice<f32>),
Cuda(CudaStorage),
}
pub(crate) trait UnaryOp {
@ -100,14 +100,14 @@ impl Storage {
pub fn device(&self) -> Device {
match self {
Self::Cpu(_) => Device::Cpu,
Self::Cuda(slice) => Device::Cuda(slice.device()),
Self::Cuda(storage) => Device::Cuda(storage.device()),
}
}
pub fn dtype(&self) -> DType {
match self {
Self::Cpu(storage) => storage.dtype(),
Self::Cuda { .. } => todo!(),
Self::Cuda(storage) => storage.dtype(),
}
}