glenn-jocher
commited on
Commit
•
99de551
1
Parent(s):
2d45de6
pt model to cpu on TF export
Browse files
export.py
CHANGED
@@ -494,7 +494,7 @@ def run(data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path'
|
|
494 |
if int8 or edgetpu: # TFLite --int8 bug https://github.com/ultralytics/yolov5/issues/5707
|
495 |
check_requirements(('flatbuffers==1.12',)) # required before `import tensorflow`
|
496 |
assert not (tflite and tfjs), 'TFLite and TF.js models must be exported separately, please pass only one type.'
|
497 |
-
model, f[5] = export_saved_model(model, im, file, dynamic, tf_nms=nms or agnostic_nms or tfjs,
|
498 |
agnostic_nms=agnostic_nms or tfjs, topk_per_class=topk_per_class,
|
499 |
topk_all=topk_all, conf_thres=conf_thres, iou_thres=iou_thres) # keras model
|
500 |
if pb or tfjs: # pb prerequisite to tfjs
|
|
|
494 |
if int8 or edgetpu: # TFLite --int8 bug https://github.com/ultralytics/yolov5/issues/5707
|
495 |
check_requirements(('flatbuffers==1.12',)) # required before `import tensorflow`
|
496 |
assert not (tflite and tfjs), 'TFLite and TF.js models must be exported separately, please pass only one type.'
|
497 |
+
model, f[5] = export_saved_model(model.cpu(), im, file, dynamic, tf_nms=nms or agnostic_nms or tfjs,
|
498 |
agnostic_nms=agnostic_nms or tfjs, topk_per_class=topk_per_class,
|
499 |
topk_all=topk_all, conf_thres=conf_thres, iou_thres=iou_thres) # keras model
|
500 |
if pb or tfjs: # pb prerequisite to tfjs
|