Skip to content

Commit

Permalink
Check CoreML models.train() mode (ultralytics#3262)
Browse files Browse the repository at this point in the history
* Check CoreML models.train() mode

* Update export.py
  • Loading branch information
glenn-jocher committed May 20, 2021
1 parent f417c69 commit 1ea05f9
Showing 1 changed file with 1 addition and 0 deletions.
1 change: 1 addition & 0 deletions models/export.py
Original file line number Diff line number Diff line change
Expand Up @@ -131,6 +131,7 @@
import coremltools as ct

print(f'{prefix} starting export with coremltools {ct.__version__}...')
assert opt.train, 'CoreML exports should be placed in model.train() mode with `python export.py --train`'
model = ct.convert(ts, inputs=[ct.ImageType('image', shape=img.shape, scale=1 / 255.0, bias=[0, 0, 0])])
f = opt.weights.replace('.pt', '.mlmodel') # filename
model.save(f)
Expand Down

0 comments on commit 1ea05f9

Please sign in to comment.