Python builtin forward is currently not supported in Torchscript

maksad: <p>I have been puzzled by the inner workings of <code>torch.jit.script</code> and how it compiles the model code. My issue is to use a Python <code>nn.Module</code> class that calls C++/cuda functions during <code>forward</code>. For detailed overview I have created this Github repository: <a href="https://github.com/maksad/torchscript-debug" class="inline-onebox" rel="noopener nofollow ugc">GitHub - maksad/torchscript-debug</a>.</p>
<p>To clarify here is a simple code:</p>
<pre><code class="lang-auto">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
</code></pre>
<p>and we can compile it like:</p>
<pre><code class="lang-auto">with torch.no_grad():
    my_nn = MyClass()
    my_nn = torch.jit.script(my_nn)
</code></pre>
<p>The <code>CARAFEPack</code> is simply <code>nn.Module</code> class, that calls <a href="https://github.com/maksad/CARAFE/blob/without_autograd/carafe/carafe.py#L109" rel="noopener nofollow ugc"><code>carafe_ext.forward</code></a>, which is a compiled <a href="https://github.com/maksad/CARAFE/tree/without_autograd/carafe/src" rel="noopener nofollow ugc">C++/cuda code</a>. <code>carafe_ext</code> is compiled to <code>.so</code> <a href="https://github.com/maksad/CARAFE/blob/without_autograd/setup.py#L16-L26" rel="noopener nofollow ugc">during installation with a setuptool</a>.</p>
<p>My expectation would be <code>torch.jit.script</code> wouldn’t have any problem to compile it, but in fact it raises <code>RuntimeError</code>:</p>
<pre><code class="lang-auto">RuntimeError:
Python builtin &lt;built-in method forward of PyCapsule object at 0x7fa48f8c4ed0&gt; 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,
            ~~~~~~~~~~~~~~~~~~ &lt;--- 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)
        ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ &lt;--- 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)
        ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ &lt;--- HERE
        return x
</code></pre>
<p>Am I missing specifics of <code>torch.jit.script</code>?</p>

eellison: <p>Hi, take a look at <a href="https://pytorch.org/tutorials/advanced/torch_script_custom_ops.html" class="inline-onebox" rel="noopener nofollow ugc">Extending TorchScript with Custom C++ Operators — PyTorch Tutorials 1.7.1 documentation</a>, that should answer your question.</p>