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 support for cuda streams. (#2532)

Browse Source
This commit is contained in:
Laurent Mazare
2024-10-02 21:30:58 +02:00
committed by GitHub
parent 936300678d
commit 7b60bda4ed
3 changed files with 24 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
candle-core/src/device.rs
View File

@ -130,6 +130,10 @@ impl Device {
Ok(Self::Cuda(crate::CudaDevice::new(ordinal)?))
}
pub fn new_cuda_with_stream(ordinal: usize) -> Result<Self> {
Ok(Self::Cuda(crate::CudaDevice::new_with_stream(ordinal)?))
}
pub fn new_metal(ordinal: usize) -> Result<Self> {
Ok(Self::Metal(crate::MetalDevice::new(ordinal)?))
}