rockstarr commited on
Commit
4c7b2bd
1 Parent(s): c55e674

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

Browse files

* fix path error in export.py

* Update export.py

updated!

* Update export.py

oops forget something

Files changed (1) hide show
  1. export.py +1 -1
export.py CHANGED
@@ -71,7 +71,7 @@ def export_torchscript(model, im, file, optimize, prefix=colorstr('TorchScript:'
71
  ts = torch.jit.trace(model, im, strict=False)
72
  d = {"shape": im.shape, "stride": int(max(model.stride)), "names": model.names}
73
  extra_files = {'config.txt': json.dumps(d)} # torch._C.ExtraFilesMap()
74
- (optimize_for_mobile(ts) if optimize else ts).save(f, _extra_files=extra_files)
75
 
76
  LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
77
  except Exception as e:
 
71
  ts = torch.jit.trace(model, im, strict=False)
72
  d = {"shape": im.shape, "stride": int(max(model.stride)), "names": model.names}
73
  extra_files = {'config.txt': json.dumps(d)} # torch._C.ExtraFilesMap()
74
+ (optimize_for_mobile(ts) if optimize else ts).save(str(f), _extra_files=extra_files)
75
 
76
  LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
77
  except Exception as e: