On this page
CUDAPluggableAllocator
class torch.cuda.CUDAPluggableAllocator(path_to_so_file, alloc_fn_name, free_fn_name)
[source]-
CUDA memory allocator loaded from a so file.
Memory allocators are compiled in .so files and loaded dynamically using ctypes. To change the active allocator use the
torch.memory.cuda.change_current_allocator()
function.- Parameters
-
- path_to_so_file (str) – Path in the filesystem to the
.so
file containing the allocator functions - alloc_fn_name (str) – Name of the function to perform the memory allocation in the so file. The signature must be: void* alloc_fn_name(ssize_t size, int device, cudaStream_t stream);
- free_fn_name (str) – Name of the function to perform the memory release in the so file. The signature must be: void free_fn_name(void* ptr, size_t size, cudaStream_t stream);
- path_to_so_file (str) – Path in the filesystem to the
Warning
This is currently supported only in unix OSs
Note
See Memory management for details on creating and using a custom allocator
© 2024, PyTorch Contributors
PyTorch has a BSD-style license, as found in the LICENSE file.
https://pytorch.org/docs/2.1/generated/torch.cuda.CUDAPluggableAllocator.html