[fix] pass input/output shapes/types to backends

jokap11 · Jan 26, 2024 · d453008 · d453008
1 parent 2be3339
commit d453008
Showing 1 changed file with 11 additions and 1 deletion.
diff --git a/mlonmcu/session/run.py b/mlonmcu/session/run.py
@@ -903,7 +903,17 @@ def _build():
                 model_artifact = model_artifact[0]
                 if not model_artifact.exported:
                     model_artifact.export(self.dir)
-                self.backend.load_model(model=model_artifact.path)
+                input_shapes = self.model.input_shapes
+                output_shapes = self.model.output_shapes
+                input_types = self.model.input_types
+                output_types = self.model.output_types
+                self.backend.load_model(
+                    model=model_artifact.path,
+                    input_shapes=input_shapes,
+                    output_shapes=output_shapes,
+                    input_types=input_types,
+                    output_types=output_types,
+                )
                 _build()
 
         self.sub_names.extend(self.artifacts_per_stage[RunStage.BUILD])