Unverified 提交 4c7b2bdc authored 作者: rockstarr's avatar rockstarr 提交者: GitHub

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

* fix path error in export.py * Update export.py updated! * Update export.py oops forget something
上级 c55e674a
......@@ -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:
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论