donothing_cuda.cu

void donothing_cuda(int alloc_size, int device_index, int deallocate) {
    cudaSetDevice((int)device_index);
    
    void *buffer;
    cudaMalloc((void**)&buffer, alloc_size);
    if (deallocate) {
        cudaFree(buffer);
    }
}

donothing_wrapper.cpp

#include <torch/extension.h>

// CUDA forward declaration
void donothing_cuda(int alloc_size, int device_index, int deallocate);

// C++ interface
void donothing(int alloc_size, int device_index, int deallocate) {
    donothing_cuda(alloc_size, device_index, deallocate);
}

// Binding to Python
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("call", &donothing, "Does exactly nothing (ok, it allocates memory)");
}

example.ipynb