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

Add the fill kernel and use it for 'ones'.

Browse Source
This commit is contained in:
laurent
2023-06-22 08:33:32 +01:00
parent fc26bab3ed
commit 0a758ffa05
3 changed files with 56 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
src/device.rs
View File

@ -82,10 +82,7 @@ impl Device {
Ok(Storage::Cpu(storage))
}
Device::Cuda(device) => {
// TODO: Instead of allocating memory on the host and transfering it,
// allocate some zeros on the device and use a shader to set them to 1.
let storage = CpuStorage::ones_impl(shape, dtype);
let storage = device.cuda_from_cpu_storage(&storage)?;
let storage = device.ones_impl(shape, dtype)?;
Ok(Storage::Cuda(storage))
}
}