Fix TorchScript on mobile export (#6183)

* fix export of TorchScript on mobile * Cleanup Co-authored-by: yinrong <[email protected]> Co-authored-by: Glenn Jocher <[email protected]>
ultralytics · Jan 4, 2022 · fb83929 · 3097254259 · Jan 22, 2022 · fb83929
1 parent 5bd6a97
commit fb83929
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/export.py b/export.py
@@ -75,7 +75,10 @@ def export_torchscript(model, im, file, optimize, prefix=colorstr('TorchScript:'
         ts = torch.jit.trace(model, im, strict=False)
         d = {"shape": im.shape, "stride": int(max(model.stride)), "names": model.names}
         extra_files = {'config.txt': json.dumps(d)}  # torch._C.ExtraFilesMap()
-        (optimize_for_mobile(ts) if optimize else ts).save(str(f), _extra_files=extra_files)
+        if optimize:  # https://pytorch.org/tutorials/recipes/mobile_interpreter.html
+            optimize_for_mobile(ts)._save_for_lite_interpreter(str(f), _extra_files=extra_files)
+        else:
+            ts.save(str(f), _extra_files=extra_files)
 
         LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
     except Exception as e: