Do not optimize CoreML TorchScript model (ultralytics#3055)

fhkiel-mlaip · May 6, 2021 · e3ce0f7 · e3ce0f7
1 parent e59979a
commit e3ce0f7
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/models/export.py b/models/export.py
@@ -75,8 +75,7 @@
         print(f'\n{prefix} starting export with torch {torch.__version__}...')
         f = opt.weights.replace('.pt', '.torchscript.pt')  # filename
         ts = torch.jit.trace(model, img, strict=False)
-        ts = optimize_for_mobile(ts)  # https://pytorch.org/tutorials/recipes/script_optimized.html
-        ts.save(f)
+        optimize_for_mobile(ts).save(f)  # https://pytorch.org/tutorials/recipes/script_optimized.html
         print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
     except Exception as e:
         print(f'{prefix} export failure: {e}')