tch-rs
tch-rs copied to clipboard
support empty_cuda_cache
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();
}
}