Check CoreML models.train() mode (#3262)

* Check CoreML models.train() mode

* Update export.py
This commit is contained in:
Glenn Jocher authored and GitHub committed 2021-05-20 15:22:36 +02:00
1 parent 7b863a34dc
commit 7b36e38cf8
1 file changed
+1
+1
View File
@@ -131,6 +131,7 @@ if __name__ == '__main__':
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)