[]-indexing should work for a CUDA tensor. For example:
auto options = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA, 0);
auto output = torch::rand({1, 3, 4, 5}, options);
for(int i = 0; i < img.rows; i++){
    for(int j = 0; j < img.cols; j++){
        output[0][0][i][j] = 1;
        output[0][1][i][j] = 2;
        output[0][2][i][j] = 3;
    }
}
assert(output.device().type() == torch::kCUDA);