支持从so文件中加载自定义NPU内存分配器。
接口原型：
torch_npu.npu.NPUPluggableAllocator(path_to_so_file, alloc_fn_name, free_fn_name)
此接口详情请参考《API参考》中的“torch.npu.npu.NPUPluggableAllocator”章节。
1 2 3 4 5 6 7 8 | import torch import torch_npu # Load the allocator new_alloc = torch_npu.npu.NPUPluggableAllocator('pluggable_allocator_extensions.so', 'my_malloc', 'my_free') # Swap the current allocator torch_npu.npu.memory.change_current_allocator(new_alloc) # This will allocate memory in the device using the new allocator npu_tensor = torch.zeros(10, device='npu') |
alloc_fn_name内存申请函数名必须与C/C++文件中函数名一致。
free_fn_name内存释放函数名必须与C/C++文件中函数名一致。