update export.py torchscript file extension

5 years ago · 9006b85d21
parent 41bb70b325
commit 9006b85d21
1 changed files with 2 additions and 2 deletions
--- a/models/export.py
+++ b/models/export.py
@ -31,7 +31,7 @@ if __name__ == '__main__':
    # TorchScript export
    try:
        print('\nStarting TorchScript export with torch %s...' % torch.__version__)
-        f = opt.weights.replace('.pt', '.torchscript')  # filename
+        f = opt.weights.replace('.pt', '.torchscript.pt')  # filename
        ts = torch.jit.trace(model, img)
        ts.save(f)
        print('TorchScript export success, saved as %s' % f)