frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-29 03:04:55 +08:00
Code Issues Packages Projects Releases Wiki Activity

expose host_emptyCache to python, fix a bug in freeing cudaHostRegist… (#134919)

Browse Source 
…ered memory

Fixes #ISSUE_NUMBER

Pull Request resolved: https://github.com/pytorch/pytorch/pull/134919
Approved by: https://github.com/eqy
This commit is contained in:
Natalia Gimelshein
2024-09-01 09:07:25 +00:00
committed by PyTorch MergeBot
parent caa04e0cae
commit c25b64a057
3 changed files with 26 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
torch/csrc/cuda/Module.cpp
View File

@ -542,6 +542,15 @@ PyObject* THCPModule_setMemoryFraction(PyObject* _unused, PyObject* args) {
Py_RETURN_NONE;
}
PyObject* THCPModule_hostEmptyCache(PyObject* _unused, PyObject* noargs) {
HANDLE_TH_ERRORS {
pybind11::gil_scoped_release no_gil;
at::cuda::CachingHostAllocator_emptyCache();
}
END_HANDLE_TH_ERRORS
Py_RETURN_NONE;
}
PyObject* THCPModule_emptyCache(PyObject* _unused, PyObject* noargs) {
HANDLE_TH_ERRORS {
pybind11::gil_scoped_release no_gil;
@ -1839,6 +1848,7 @@ static struct PyMethodDef _THCPModule_methods[] = {
THCPModule_cudaHostAllocator,
METH_NOARGS,
nullptr},
{"_host_emptyCache", THCPModule_hostEmptyCache, METH_NOARGS, nullptr},
{"_cuda_cudaCachingAllocator_raw_alloc",
THCPModule_cudaCachingAllocator_raw_alloc,
METH_VARARGS,