tch-rs icon indicating copy to clipboard operation
tch-rs copied to clipboard

support empty_cuda_cache

Open wenhaozhao opened this issue 5 months ago • 0 comments

Added empty_cuda_cache method to release GPU cache.

tch::torch_sys::c10_cuda::empty_cuda_cache();

Implemented the same functionality as torch.cuda.empty_cache() in pytorch.

Example

pub fn free_gpu_memory() {
    #[cfg(not(target_os = "macos"))]
    {
        use tch::torch_sys;
        torch_sys::c10_cuda::empty_cuda_cache();
    }
}

wenhaozhao avatar Sep 05 '24 06:09 wenhaozhao