If I want to call other methods in my model(not forward), what should I do? I try to use torch.jit.script_method to decorate it but I can’t call it in C++.
class Net(torch.jit.ScriptModule): def __init__(self): super(Net, self).__init__() self.conv = nn.Conv2d(1, 1, 3) @torch.jit.script_method def forward(self, x): return self.conv(x) @torch.jit.script_method def weighted_kernel_sum(self, weight): return weight * self.conv.weigh n = Net() n.save("model.pt")
What I want to say is that if I use this way to convert my model, can I call the weighted_kernel_sum method in the C++ code?
run_method work for you?
Thank you for your reply. I don’t quite understand that, could you please give me an example in detail?
You follow the tutorial and at step 4 you replace
at::Tensor output = module->forward(inputs).toTensor();
torch::Tensor output = module->run_method("weighted_kernel_sum", input).toTensor();
where input is your input tensor.
It works. Thank you for your help!