TL;DR

torch.cuda.memory.CUDAPluggableAllocatorを使え

サンプルコード

よっこらしょ。
　　　　∧＿∧　　ミ　＿　ドスッ
　　　（　　　　）┌─┴┴─┐
　　　 /　　　　つ.　終　　了　|
　　 :/o　　 /´ .└─┬┬─┘
　　（＿（＿）　；;､;｡;|　|
　　この記事は無事に終了しました
　　ありがとうございました

流石に少し解説
公式ドキュメントを読む,

class torch.cuda.CUDAPluggableAllocator(path_to_so_file, alloc_fn_name, free_fn_name)

CUDA memory allocator loaded from a so file.とある。要は共有オブジェクトファイルを読み込んでallocate関数とfree関数の名前をかけとのこと, ここで気になるのはallocate関数とfree関数がどういう引数を取るかだが, それはソースコードを読めばわかる。

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

というわけで, allocate関数は(ssize_t size, int device, cudaStream_t stream)を引数にとり, free関数は(void* ptr, size_t size, cudaStream_t stream)を引数にとれば良い(CUDAの構造体に関しての説明は省略)

実装してみる

#include <cuda_runtime_api.h>
#include <stdlib.h>

extern "C" {

void *custom_malloc(ssize_t size, int device, cudaStream_t stream) {
    void *ptr;
    #ifdef MANAGED
    cudaMallocManaged(&ptr, size);
    #elif CPU
    cudaMallocHost(&ptr, size);
    #else
    cudaMallocAsync(&ptr, size, stream);
    #endif
    return ptr;
}

void custom_free(void *ptr, ssize_t size, int device, cudaStream_t stream) {
    #ifdef MANAGED
    cudaFree(ptr);
    #elif CPU
    cudaFreeHost(ptr);
    #else
    cudaFreeAsync(ptr, stream);
    #endif
}

}

こいつをPythonから呼び出すには, 共有オブジェクトとしてコンパイルして

# memory allocator 変更
new_alloc = torch.cuda.memory.CUDAPluggableAllocator('allocator.so', 'custom_malloc', 'custom_free')
torch.cuda.memory.change_current_allocator(new_alloc)

こうすれば良い

よっこらしょ。
　　　　∧＿∧　　ミ　＿　ドスッ
　　　（　　　　）┌─┴┴─┐
　　　 /　　　　つ.　終　　了　|
　　 :/o　　 /´ .└─┬┬─┘
　　（＿（＿）　；;､;｡;|　|
　　この記事は本当に無事に終了しました
　　ありがとうございました