I wrote a custom module extension in C++, and compiled it with Pybind. However, when I load it in Python i get a missing symbol error:
_ZNK5torch8autograd4Node4nameB5cxx11Ev
// Inherit from Function
class LinearFunction : public Function<LinearFunction> {
public:
// Note that both forward and backward are static functions
// bias is an optional argument
static torch::Tensor forward(
AutogradContext *ctx, torch::Tensor input, torch::Tensor weight, torch::Tensor bias = torch::Tensor()) {
ctx->save_for_backward({input, weight, bias});
auto output = input.mm(weight.t());
if (bias.defined()) {
output += bias.unsqueeze(0).expand_as(output);
}
return output;
}
static tensor_list backward(AutogradContext *ctx, tensor_list grad_outputs) {
auto saved = ctx->get_saved_variables();
auto input = saved[0];
auto weight = saved[1];
auto bias = saved[2];
auto grad_output = grad_outputs[0];
auto grad_input = grad_output.mm(weight);
auto grad_weight = grad_output.t().mm(input);
auto grad_bias = torch::Tensor();
if (bias.defined()) {
grad_bias = grad_output.sum(0);
}
return {grad_input, grad_weight, grad_bias};
}
};
void testing(torch::Tensor x, torch::Tensor weight, torch::Tensor bias)
{
auto y = LinearFunction::apply(x, weight);
y.sum().backward();
}
PYBIND11_MODULE(libModule, m)
{
m.def("testing", &testing, "testing");
}
The error only occurs if I have LinearFunction::apply(x, weight); not commented out.
I have the following modules I am linking against:
"lib/libtorch.so",
"lib/libtorch_cpu.so",
"lib/libc10.so",
"lib/libgomp-75eea7e8.so.1",
"lib/libtorch_python.so",
"lib/libshm.so",
Everything works fine if i comment out LinearFunction::apply but if I enable it I get that missing symbol error _ZNK5torch8autograd4Node4nameB5cxx11Ev