I don’t know if it will solve, but you should maybe cast the input and output in the declaration of pybind11, like it is done here: https://github.com/LeviViana/torch_sampling/blob/master/csrc/sampling.cpp#L44 2
BTW, this repo contains an implementation of a pytorch extension, it can serve as tutorial as well.