Exporting yolov4-tiny to torchscript failing

[Harshil Gupta]

I'm trying to export a trained yolov4 tiny model to torchscript. The goal is to use that torchscript with aws's elastic inference.
Tried the two export methods.

jit.script is failing with this error:

Module 'YOLOLayer' has no attribute 'layers' (This attribute exists on the Python module, but we failed to convert Python type: 'list' to a TorchScript type.): File "/home/ubuntu/cv_tasks/tracking/yolov4/models_yolov4/models.py", line 391 if ASFF: i, n = self.index, self.nl # index in layers, number of layers p = out[self.layers[i]] ~~~~~~~~~~~ <--- HERE bs, _, ny, nx = p.shape # bs, 255, 13, 13 if (self.nx, self.ny) != (nx, ny):

and jit.trace is failing with
torch.jit.TracingCheckError: Tracing failed sanity checks! ERROR: Graphs differed across invocations! Graph diff: graph(%self.1 : __torch__.models_yolov4.models.Darknet,

I tried check_trace=False in jit.trace to circumvent the above error. A scripted model was created.
But when loading the script using torch 1.5.1 and torcheia gave the following error:

File "updated_inference.py", line 177, in load_yolov4
eia_model = torcheia.jit.attach_eia(model, 0)
File "/home/ubuntu/.venv/lib/python3.6/site-packages/torcheia/jit/init.py", line 5, in attach_eia
newobj = _torch_eia._tag_eia_device(obj._c, device)
RuntimeError: class 'torch.models_yolov4.models.___torch_mangle_93.YOLOLayer' already defined.

Please help.