mirrors
/
yolov5
mirror of https://github.com/ultralytics/yolov5.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

TorchScript `torch==1.7.0` Path support (#5781) 

* fix path error in export.py

* Update export.py

updated!

* Update export.py

oops forget something
pull/5804/head
rockstarr 2021-11-25 18:18:30 +08:00 committed by GitHub
parent c55e674ac3
commit 4c7b2bdc30
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
export.py
View File

@ -71,7 +71,7 @@ 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(f, _extra_files=extra_files)
(optimize_for_mobile(ts) if optimize 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: