diff --git a/mlonmcu/session/run.py b/mlonmcu/session/run.py
index c7ac4cbb4..e71a4f320 100644
--- 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])