Consistent saved_model output format (ultralytics#7032)

fhkiel-mlaip · Mar 18, 2022 · 0d5433a · 0d5433a
1 parent 4b81e09
commit 0d5433a
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/export.py b/export.py
@@ -275,7 +275,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), [spec])
+            tfm.__call__ = tf.function(lambda x: frozen_func(x)[0], [spec])
             tfm.__call__(im)
             tf.saved_model.save(
                 tfm,

diff --git a/models/common.py b/models/common.py
@@ -441,7 +441,7 @@ def forward(self, im, augment=False, visualize=False, val=False):
         else:  # TensorFlow (SavedModel, GraphDef, Lite, Edge TPU)
             im = im.permute(0, 2, 3, 1).cpu().numpy()  # torch BCHW to numpy BHWC shape(1,320,192,3)
             if self.saved_model:  # SavedModel
-                y = (self.model(im, training=False) if self.keras else self.model(im)[0]).numpy()
+                y = (self.model(im, training=False) if self.keras else self.model(im)).numpy()
             elif self.pb:  # GraphDef
                 y = self.frozen_func(x=self.tf.constant(im)).numpy()
             else:  # Lite or Edge TPU