AutoUpdate TensorFlow in export.py (#9447)

Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com> Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
ultralytics · Sep 16, 2022 · db06f49 · db06f49
1 parent fe10b4a
commit db06f49
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/export.py b/export.py
@@ -309,6 +309,7 @@ def export_saved_model(model,
                        keras=False,
                        prefix=colorstr('TensorFlow SavedModel:')):
     # YOLOv5 TensorFlow SavedModel export
+    check_requirements('tensorflow' if torch.cuda.is_available() else 'tensorflow-cpu')
     import tensorflow as tf
     from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2