gistfile1.txt

import torch
from torch.utils.cpp_extension import load_inline

cpp_code = """
torch::Tensor to_gray(torch::Tensor input);

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("to_gray", &to_gray, "Convert RGB to Grayscale (CUDA)");
}
"""

cuda_kernel_code = """
torch::Tensor to_gray(torch::Tensor input) {
  auto output = torch::empty({input.size(0), input.size(1)}, input.options());  
  return output;
}
"""

cuda_module = load_inline(
    name="to_gray_cuda",
    cpp_sources=cpp_code, 
    cuda_sources=cuda_kernel_code, 
    functions=["to_gray"],
    with_cuda=True,
    verbose=True,
    extra_cflags=["-std=c++17", "-ftime-report"],
    extra_cuda_cflags=["-arch=sm_90"]  

)