I have been puzzled by the inner workings of torch.jit.script and how it compiles the model code. My issue is to use a Python nn.Module class that calls C++/cuda functions during forward. For detailed overview I have created this Github repository: GitHub - maksad/torchscript-debug 1.
To clarify here is a simple code:
import torch.nn as nn
import torch
from carafe import CARAFEPack

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

        self.conv = nn.Conv2d(in_channels=1, out_channels=1, kernel_size=2)
        self.upsampling = CARAFEPack(channels=1, scale_factor=2)

    def forward(self, x):
        x = self.conv(x)
        x = self.upsampling(x)
        return x

and we can compile it like:
with torch.no_grad():
    my_nn = MyClass()
    my_nn = torch.jit.script(my_nn)

The CARAFEPack is simply nn.Module class, that calls carafe_ext.forward, which is a compiled C++/cuda code. carafe_ext is compiled to .so during installation with a setuptool.
My expectation would be torch.jit.script wouldn’t have any problem to compile it, but in fact it raises RuntimeError:
RuntimeError:
Python builtin <built-in method forward of PyCapsule object at 0x7fa48f8c4ed0> is currently not supported in Torchscript:
  File "/home/ubuntu/.virtualenvs/temp/lib/python3.6/site-packages/carafe/carafe.py", line 109
        rmasks = masks.new_zeros(masks.size())
        if features.is_cuda:
            carafe_ext.forward(features, rfeatures, masks, rmasks, self.up_kernel,
            ~~~~~~~~~~~~~~~~~~ <--- HERE
                            self.up_group, self.scale_factor, routput, output)
        else:
'CARAFEPack.forward_carafe' is being compiled since it was called from 'CARAFEPack.feature_reassemble'
  File "/home/ubuntu/.virtualenvs/temp/lib/python3.6/site-packages/carafe/carafe.py", line 92
    def feature_reassemble(self, x: Tensor, mask: Tensor):
        x = self.forward_carafe(x, mask, self.up_kernel, self.up_group, self.scale_factor)
        ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
        return x
'CARAFEPack.feature_reassemble' is being compiled since it was called from 'CARAFEPack.forward'
  File "/home/ubuntu/.virtualenvs/temp/lib/python3.6/site-packages/carafe/carafe.py", line 121
        mask = self.kernel_normalizer(mask)
        x = self.feature_reassemble(x, mask)
        ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
        return x

Am I missing specifics of torch.jit.script?