huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-17 02:58:50 +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

7
src/tensor.rs
View File

@ -250,7 +250,12 @@ impl Tensor {
let data = S::cpu_storage_as_slice(cpu_storage)?;
Ok(self.strided_index().map(|i| data[i]).collect())
}
Storage::Cuda(_) => todo!(),
Storage::Cuda(slice) => {
// TODO: Would it be possible to only fetch the necessary data?
let cpu_storage = slice.to_cpu_storage()?;
let data = S::cpu_storage_as_slice(&cpu_storage)?;
Ok(self.strided_index().map(|i| data[i]).collect())
}
}
}