pytorch / 2 / generated / torch.cuda.cudapluggableallocator.html

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

© 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