From 9006b85d21e33d2ca5390fa6ab469810380e5c9c Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Fri, 10 Jul 2020 11:56:01 -0700 Subject: [PATCH] update export.py torchscript file extension --- models/export.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/models/export.py b/models/export.py index 990c86e..2097df5 100644 --- 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) @@ -62,7 +62,7 @@ if __name__ == '__main__': print('\nStarting CoreML export with coremltools %s...' % ct.__version__) # convert model from torchscript and apply pixel scaling as per detect.py - model = ct.convert(ts, inputs=[ct.ImageType(name='images', shape=img.shape, scale=1/255.0, bias=[0, 0, 0])]) + model = ct.convert(ts, inputs=[ct.ImageType(name='images', shape=img.shape, scale=1 / 255.0, bias=[0, 0, 0])]) f = opt.weights.replace('.pt', '.mlmodel') # filename model.save(f) print('CoreML export success, saved as %s' % f)