CUDAPluggableAllocator — PyTorch 2.0 documentation (original) (raw)

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);

Warning

This is currently supported only in unix OSs