Only tuples, lists and Variables are supported as JIT inputs/outputs. Dictionaries and strings are also accepted, but their usage is not recommended. Here, received an input of unsupported type: int

username: Kracozebr
cooked: <p>Hi I’m trying to export PyTorch custom layer to  to onnx and got this error:</p>
<pre><code class="lang-auto">Only tuples, lists and Variables are supported as JIT inputs/outputs. Dictionaries and strings are also accepted, but their usage is not recommended. Here, received an input of unsupported type: int
</code></pre>
<p>I checked same issues and tried to resolve:</p><aside class="quote quote-modified" data-post="1" data-topic="118199">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ashraf_g/48/37112_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/onnx-export-with-multiple-arguments-only-tuples-lists-and-variables-are-supported-as-jit-inputs-outputs/118199">ONNX Export with Multiple Arguments: Only tuples, lists and Variables are supported as JIT inputs/outputs</a> <a class="badge-category__wrapper " href="/c/jit/13"><span data-category-id="13" data-drop-close="true" class="badge-category " title="A category for TorchScript and the PyTorch JIT compiler"><span class="badge-category__name">jit</span></span></a>
  </div>
  <blockquote>
    Hi, everyone! 

<a name="issue-summary" class="anchor" href="#issue-summary"></a>Issue Summary:
I’m trying to convert the ESM-1b protein transformer model from PyTorch to ONNX. I run into an issue when I want to provide an extra argument to the model. 

<a name="conversion-steps-code-and-error-message" class="anchor" href="#conversion-steps-code-and-error-message"></a>Conversion Steps, Code, and Error Message:
Here is my conversion script (named convert_onnx_esm.py): 
import os
import torch
import torch.onnx
import argparse
from esm.pretrained import load_model_and_alphabet_local


parser = argparse.ArgumentParser()

parser.add_argument("--model-path", type=str, required=Tr…
  </blockquote>
</aside>
<aside class="quote quote-modified" data-post="1" data-topic="88566">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/sayak_paul/48/12743_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/runtimeerror-only-tuples-lists-and-variables-supported-as-jit-inputs-outputs-dictionaries-and-strings-are-also-accepted-but-their-usage-is-not-recommended-but-got-unsupported-type-nonetype/88566">RuntimeError: Only tuples, lists and Variables supported as JIT inputs/outputs. Dictionaries and strings are also accepted but their usage is not recommended. But got unsupported type NoneType</a> <a class="badge-category__wrapper " href="/c/vision/5"><span data-category-id="5" data-drop-close="true" class="badge-category " title="Topics related to either pytorch/vision or vision research related topics"><span class="badge-category__name">vision</span></span></a>
  </div>
  <blockquote>
    Hi. 
I am trying to convert the <a href="https://github.com/NVIDIA/FastPhotoStyle" rel="noopener nofollow ugc">FastPhotoStyle model (PyTorch) by NVIDIA</a>. The stylization module takes four inputs i.e. the style image, the content image, and two segmentation masks. 
I am running the conversion step like so - 
torch.onnx.export(stylization_module, [cont_img, styl_img, cont_img, styl_img], 'stylization_module.onnx', input_names=['content', 'style', 'content_seg', 'style_seg'], output_names=['stylized_image'])`

It is producing: 
RuntimeError: Only tuples, lists and Variables su…
  </blockquote>
</aside>
<p>
But cannot do it.<br>
There is my code:</p>
<pre><code class="lang-auto">import torch
import dcn_op_v2
from torch.autograd import Function
from torch.nn import Module
import torch.nn as nn
import math


class DeformableConv2DFunction(Function):
    @staticmethod
    def forward(ctx, input_tensor, weight, bias, offset, mask, stride, pad, dilation, deformable_groups):
        ctx.stride_h = stride[0]
        ctx.stride_w = stride[1]
        ctx.pad_h = pad[0]
        ctx.pad_w = pad[1]
        ctx.dilation_h = dilation[0]
        ctx.dilation_w = dilation[1]
        ctx.deformable_groups = deformable_groups

        output = dcn_op_v2.forward(
            input_tensor,
            weight,
            bias,
            offset,
            mask,
            ctx.stride_h, ctx.stride_w,
            ctx.pad_h, ctx.pad_w,
            ctx.dilation_h, ctx.dilation_w,
            ctx.deformable_groups
        )
        ctx.save_for_backward(input_tensor, weight, offset, mask, bias)
        return output
    
    @staticmethod
    def backward(ctx, *grad_outputs):
        input_tensor, weight, offset, mask, bias = ctx.saved_tensors
        grad_input, grad_weight, grad_bias, grad_offset, grad_mask = dcn_op_v2.backward(
            input_tensor,
            weight,
            bias,
            offset,
            mask,
            grad_outputs[0],
            ctx.stride_h, ctx.stride_w,
            ctx.pad_h, ctx.pad_w,
            ctx.dilation_h, ctx.dilation_w,
            ctx.deformable_groups
        )
        
        return grad_input, grad_weight, grad_bias, grad_offset, grad_mask, \
            None, None, None, None

class DeformableConv2DLayer(Module):
    def __init__(self,
                 in_channels,
                 out_channels,
                 kernel_size,
                 stride,
                 padding,
                 dilation,
                 deformable_groups):
        super().__init__()
        def typexam(x):
            if type(x)==int:
                return (x, x)
            elif type(x)==tuple and len(x)==2:
                return x
            else:
                raise TypeError
        kernel_size = typexam(kernel_size)
        stride = typexam(stride)
        padding = typexam(padding)
        dilation = typexam(dilation)
        self.stride = stride
        self.pad = padding
        self.dilation = dilation
        self.deformable_groups = deformable_groups
        self.weight = nn.Parameter(torch.Tensor(out_channels, in_channels, *kernel_size))
        self.bias = nn.Parameter(torch.Tensor(out_channels))
        self.conv_offset_mask = nn.Conv2d(in_channels,
                                          self.deformable_groups * 3 * kernel_size[0] * kernel_size[1],
                                          kernel_size=kernel_size,
                                          stride=self.stride,
                                          padding=self.pad,
                                          bias=True)
        self.reset_parameters(in_channels, kernel_size)
        self.init_offset()

    def init_offset(self):
        self.conv_offset_mask.weight.data.zero_()
        self.conv_offset_mask.bias.data.zero_()

    def reset_parameters(self, in_channels, kernel_size):
        n = in_channels
        for k in kernel_size:
            n *= k
        stdv = 1. / math.sqrt(n)
        self.weight.data.uniform_(-stdv, stdv)
        self.bias.data.zero_()

    def forward(self, inputs):
        out = self.conv_offset_mask(inputs)
        o1, o2, mask = torch.chunk(out, 3, dim=1)
        offset = torch.cat((o1, o2), dim=1)
        mask = torch.sigmoid(mask)
        return DeformableConv2DFunction.apply(inputs, self.weight, self.bias, offset, mask, self.stride, self.pad, self.dilation, self.deformable_groups)

from torch.onnx import register_custom_op_symbolic

def dcn_v2(g, self):
    return g.op("ai.onnx.contrib::Dcn_v2", self)

register_custom_op_symbolic('::dcn_v2', dcn_v2, 1)


import io
import onnx


N, inC, inH, inW = 2, 2, 4, 4
outC = 2
kH, kW = 3, 3
stride = (1,1)
padding = (1,1)
dilation= (1,1)
deformable_groups = 1

t = torch.randn(N, inC, inH, inW)

offset = torch.randn(N, deformable_groups * 2 * kW * kH, inH, inW)

mask = torch.rand(N, deformable_groups * 1 * kW * kH, inH, inW)
# mask.data.zero_()
mask = torch.sigmoid(mask)

weight = torch.randn(outC, inC, kH, kW)
bias = torch.rand(outC)



# Export model to ONNX
t_model = DeformableConv2DLayer(inC, outC, (kH,kW), stride, padding, dilation, deformable_groups)
input_nn = (t, weight, bias, offset, mask, stride, padding, dilation, torch.tensor(deformable_groups, dtype=torch.int32))
with torch.no_grad():
    torch.onnx.export(t_model, input_nn, 'models/dcn.onnx',  export_params=True)
</code></pre>
<p>Full ouput of error is:</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "custom_onnx_op.py", line 158, in &lt;module&gt;
    torch.onnx.export(t_model, input_nn, 'models/dcn.onnx',  export_params=True)
  File "/home/appuser/.local/lib/python3.6/site-packages/torch/onnx/__init__.py", line 276, in export
    custom_opsets, enable_onnx_checker, use_external_data_format)
  File "/home/appuser/.local/lib/python3.6/site-packages/torch/onnx/utils.py", line 94, in export
    use_external_data_format=use_external_data_format)
  File "/home/appuser/.local/lib/python3.6/site-packages/torch/onnx/utils.py", line 701, in _export
    dynamic_axes=dynamic_axes)
  File "/home/appuser/.local/lib/python3.6/site-packages/torch/onnx/utils.py", line 459, in _model_to_graph
    use_new_jit_passes)
  File "/home/appuser/.local/lib/python3.6/site-packages/torch/onnx/utils.py", line 420, in _create_jit_graph
    graph, torch_out = _trace_and_get_graph_from_model(model, args)
  File "/home/appuser/.local/lib/python3.6/site-packages/torch/onnx/utils.py", line 380, in _trace_and_get_graph_from_model
    torch.jit._get_trace_graph(model, args, strict=False, _force_outplace=False, _return_inputs_states=True)
  File "/home/appuser/.local/lib/python3.6/site-packages/torch/jit/_trace.py", line 1139, in _get_trace_graph
    outs = ONNXTracedModule(f, strict, _force_outplace, return_inputs, _return_inputs_states)(*args, **kwargs)
  File "/home/appuser/.local/lib/python3.6/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/appuser/.local/lib/python3.6/site-packages/torch/jit/_trace.py", line 93, in forward
    in_vars, in_desc = _flatten(args)
RuntimeError: Only tuples, lists and Variables are supported as JIT inputs/outputs. Dictionaries and strings are also accepted, but their usage is not recommended. Here, received an input of unsupported type: int

</code></pre>
<p>So I can not understand how to deal with it.</p>