mirror of
https://github.com/ultralytics/yolov5.git
synced 2025-06-03 14:49:29 +08:00
update export.py torchscript file extension
This commit is contained in:
parent
41bb70b325
commit
9006b85d21
@ -31,7 +31,7 @@ if __name__ == '__main__':
|
|||||||
# TorchScript export
|
# TorchScript export
|
||||||
try:
|
try:
|
||||||
print('\nStarting TorchScript export with torch %s...' % torch.__version__)
|
print('\nStarting TorchScript export with torch %s...' % torch.__version__)
|
||||||
f = opt.weights.replace('.pt', '.torchscript') # filename
|
f = opt.weights.replace('.pt', '.torchscript.pt') # filename
|
||||||
ts = torch.jit.trace(model, img)
|
ts = torch.jit.trace(model, img)
|
||||||
ts.save(f)
|
ts.save(f)
|
||||||
print('TorchScript export success, saved as %s' % f)
|
print('TorchScript export success, saved as %s' % f)
|
||||||
|
Loading…
x
Reference in New Issue
Block a user