Call custom functions by jit in C++

jit
1

I have created the following class and saved scripted_module and loaded in C++ API by:

    torch::jit::script::Module module;
    module = torch::jit::load("scriptmodule.pt");

Now, the question is how can I call func from C++?

import torch
import torch.nn as nn

class MyModule(nn.Module):
    def __init__(self):
        super(MyModule, self).__init__()

    @torch.jit.ignore
    def get_rand(self):
        return torch.randint(0, 2, size=[1])

    def forward(self):
        pass 
    
    @torch.jit.export
    def func(self):
        done = self.get_rand()
        print (done)

scripted_module = torch.jit.script(MyModule())
1 Like
2

Does this answer help?

1 Like
3

Also just FYI @torch.jit.ignore functions cannot be exported, if you do scripted_module.save("out.pt") you’ll get

Traceback (most recent call last):
  File "../test.py", line 196, in <module>
    scripted_module.save('s.pt')
  File "/home/pytorch/torch/jit/__init__.py", line 1621, in save
    return self._c.save(*args, **kwargs)
RuntimeError: 
Could not export Python function call 'get_rand'. Remove calls to Python functions before export. Did you forget add @script or @script_method annotation? If this is a nn.ModuleList, add it to __constants__:
  File "../test.py", line 192
    @torch.jit.export
    def func(self):
        done = self.get_rand()
               ~~~~~~~~~~~~~ <--- HERE
        print (done)

Which is expected behavior (since saved models are expected to run without a Python runtime attached)

1 Like
4

Thanks, this worked.

5

Yes, I faced a similar example and as you mentioned it makes sense.