github.com/alwaysproblem/mlserving-tutorial@v0.0.0-20221124033215-121cfddbfbf4/TorchServing/CustomOp/src/add_index.cpp (about)

     1  #include <torch/script.h>  // One-stop header.
     2  
     3  #include <iostream>
     4  #include <memory>
     5  
     6  template <typename T>
     7  void add_index_kernel(T* input_data_ptr, int64_t sizes, T* output_data_ptr) {
     8    for (size_t i = 0; i < sizes; i++) {
     9      output_data_ptr[i] = input_data_ptr[i] + (T)i;
    10    }
    11  }
    12  
    13  torch::Tensor add_index(torch::Tensor input_tensor) {
    14    auto output = torch::zeros_like(input_tensor);
    15    size_t ndim = input_tensor.dim();
    16    int64_t dims = 1;
    17    for (size_t i = 0; i < ndim; ++i) {
    18      dims *= input_tensor.size(i);
    19    }
    20    AT_DISPATCH_ALL_TYPES(input_tensor.scalar_type(), "add_index", ([&] {
    21                            add_index_kernel<scalar_t>(
    22                                input_tensor.data_ptr<scalar_t>(), dims,
    23                                output.data_ptr<scalar_t>());
    24                          }));
    25    return output.clone();
    26  }
    27  
    28  TORCH_LIBRARY(my_ops, m) { m.def("add_index", &add_index); }