I was wondering if there is any way to pass a Torchscript module from python to C++, without having to go through serialization.
I have my own custom C++ library(not just an extension) wrapped with Pybind11. The goal is to pass a torchscript module without too much overhead, so ideally I want to avoid passing through the filesystem, or strings of Python code.
I have currently tried creating a C++ function which accepts torch::jit::script::Module, but when I pass a torch.ScriptModule it fails as different type. My guess is because PyTorch doesn’t export the Pybind code for Module in torch/Extension.h, or I am doing something wrong with the linking.
Ideally it should look something like this.
void ExampleFunction(torch::jit::script::Module& module){
// Do something
};
//Pybind11 code
m.def("CppExampleFunction",&ExampleFunction);
class MyModule(nn.Module):
def __init__(self):
super(MyModule, self).__init__()
def forward(self, input):
return 10 - input
my_script_module = torch.jit.script(MyModule())
CppExampleFunction(my_script_module)