Export to torchscript problem

[GoncaloR00]

I had problems exporting the model to the torchscript format using the 'pretrained/best.pth' file using torch 1.13.0. The fix is easy, but it's worth explaining anyway, as someone may need it.

The keys didn't match with the model using the export.py script, but they matched using the test_image.py script. In the test_image.py, there is an additional operation in line 38:

model = torch.nn.DataParallel(model)

Adding this line to the export.py script doesn't solve the problem because Torchscript doesn't support DataParallel.

To get the keys to match, I changed the key names of the pre-trained weights file using something like the following:

from collections import OrderedDict
import torch

loaded_dict = torch.load('pretrained/best.pth')
new_dict = OrderedDict()
for key in loaded_dict:
    new_dict[key.split('module.')[1]] = loaded_dict[key]

And then I loaded the new_dict in line 127.