Problems when exporting model by JIT tracing because of Mish

I created a Unet model using this line -

learn = unet_learner(dls, resnet18, act_cls=Mish, opt_func=ranger, n_out=2)

Then, I trained the model. Now, I am trying to export the model for inference using JIT trace -, trans(image).float().unsqueeze_(0)), '').

But it is throwing the error below -

Could not export Python function call 'MishJitAutoFn'. Remove calls to Python functions before export. Did you forget to add @script or @script_method annotation? If this is a nn.ModuleList, add it to __constants__:

As per error, function should have @script annotation. I checked in fastai code, there already exists @script.

def _mish_jit_fwd(x): return x.mul(torch.tanh(F.softplus(x)))

Can anybody tell me what’s wrong and how to solve this problem?

1 Like