I’ve removed some code to just print the thread and block indices and it seems to work fine:
from torch.utils import cpp_extension

cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

__global__ void tmp_kernel(void)
{
    const int ci = threadIdx.x;
    const int cj = blockIdx.x;
    printf("Hello from block %d, thread %d ",cj, ci);
}

void tmp_cuda(void) {
    const long threads = 5;
    const dim3 blocks(5);
    tmp_kernel<<<blocks, threads>>>();
    return;
}
"""

cpp_source = """
void tmp_cuda(void);

void tmp(void){
    tmp_cuda();
}
"""

module = torch.utils.cpp_extension.load_inline(
    name="tmp",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions="tmp",
    verbose=True,
)

module.tmp()