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

Warning

This is currently supported only in unix OSs

Note

See Memory management for details on creating and using a custom allocator