@@ -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, |
@@ -441,7 +441,7 @@ class DetectMultiBackend(nn.Module): | |||
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 |