huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-16 18:48:51 +00:00
Code Issues Packages Projects Releases Wiki Activity

Support cuda in to_vec3.

Browse Source
This commit is contained in:
laurent
2023-06-22 12:22:51 +01:00
parent 0689d62548
commit f052020ba2
1 changed files with 16 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
src/tensor.rs
View File

@ -359,8 +359,7 @@ impl Tensor {
pub fn to_vec3<S: crate::WithDType>(&self) -> Result<Vec<Vec<Vec<S>>>> {
let (dim1, dim2, dim3) = self.shape().r3()?;
match &self.storage {
Storage::Cpu(cpu_storage) => {
let from_cpu_storage = |cpu_storage: &crate::CpuStorage| {
let data = S::cpu_storage_as_slice(cpu_storage)?;
let mut top_rows = vec![];
let mut src_index = self.strided_index();
@ -374,8 +373,10 @@ impl Tensor {
}
assert!(src_index.next().is_none());
Ok(top_rows)
}
Storage::Cuda { .. } => todo!(),
};
match &self.storage {
Storage::Cpu(storage) => from_cpu_storage(storage),
Storage::Cuda(storage) => from_cpu_storage(&storage.to_cpu_storage()?),
}
}