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

* Check CoreML models.train() mode * Update export.py
Lechtr · May 20, 2021 · 7b36e38 · 7b36e38
1 parent 7b863a3
commit 7b36e38
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/models/export.py b/models/export.py
@@ -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)