glenn-jocher commited on
Commit
7b36e38
·
unverified ·
1 Parent(s): 7b863a3

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

Browse files

* Check CoreML models.train() mode

* Update export.py

Files changed (1) hide show
  1. models/export.py +1 -0
models/export.py CHANGED
@@ -131,6 +131,7 @@ if __name__ == '__main__':
131
  import coremltools as ct
132
 
133
  print(f'{prefix} starting export with coremltools {ct.__version__}...')
 
134
  model = ct.convert(ts, inputs=[ct.ImageType('image', shape=img.shape, scale=1 / 255.0, bias=[0, 0, 0])])
135
  f = opt.weights.replace('.pt', '.mlmodel') # filename
136
  model.save(f)
 
131
  import coremltools as ct
132
 
133
  print(f'{prefix} starting export with coremltools {ct.__version__}...')
134
+ assert opt.train, 'CoreML exports should be placed in model.train() mode with `python export.py --train`'
135
  model = ct.convert(ts, inputs=[ct.ImageType('image', shape=img.shape, scale=1 / 255.0, bias=[0, 0, 0])])
136
  f = opt.weights.replace('.pt', '.mlmodel') # filename
137
  model.save(f)