SavedModel TF Serve Fix

Fix for #7205 proposed by @tylertroy
ultralytics · Mar 31, 2022 · 98109b0 · 98109b0
1 parent df9008e
commit 98109b0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/export.py b/export.py
@@ -276,7 +276,7 @@ def export_saved_model(model, im, file, dynamic,
             m = m.get_concrete_function(spec)
             frozen_func = convert_variables_to_constants_v2(m)
             tfm = tf.Module()
-            tfm.__call__ = tf.function(lambda x: frozen_func(x)[0], [spec])
+            tfm.__call__ = tf.function(lambda x: frozen_func(x)[:4], [spec])
             tfm.__call__(im)
             tf.saved_model.save(
                 tfm,