Failing to create TorchScript module using torch.jit.script()

[Mark-M2L]

Hello all,

Thank you very much for this wonderful library. It is really useful. I have trained the fastai unet model with ResNet34 as backbone and try to port it to C++ at the moment (and perhaps later, convert it to torch_tensorrt). So, for the step to compile the TorchScript module, torch.jit.script() is used. However, this function returns a failure.

The code is as follows:

import torch
import torchvision
from icevision.models import *
from icevision.all import *

presize, size = 384, 224
presize, size = ImgSize(presize, int(presize*1.0)), ImgSize(size, int(size*1.0))

model_type = models.fastai.unet
backbone = model_type.backbones.resnet34()
model = model_type.model(backbone=backbone, num_classes=9, img_size=size)

scripted_model = torch.jit.script(model)

torch.jit.save(scripted_model, 'resnet_unet_224.pt')

The error obtained:

RuntimeError: Hook 'hook_fn' on module 'ReLU' was expected to only have exactly 3 inputs but it had 4 inputs. This error occured while scripting the forward hook 'hook_fn' on module ReLU. If you did not want to script this hook remove it from the original NN module before scripting. This hook was expected to have the following signature: hook_fn(self, input: Tuple[Tensor], output: Tensor). The type of the output arg is the returned type from either the forward method or the previous hook if it exists. Note that hooks can return anything, but if the hook is on a submodule the outer module is expecting the same return type as the submodule's forward.

Further, I get the same error when trying to compile the model using torch_tensorrt. Can you help with creating the TorchScript? Thanks!