2 files changed, 8 insertions, 2 deletions
diff --git a/libs/input/TfLiteMotionPredictor.cpp b/libs/input/TfLiteMotionPredictor.cpp
index 3b061d1cf1..85fa176129 100644
--- a/libs/input/TfLiteMotionPredictor.cpp
+++ b/libs/input/TfLiteMotionPredictor.cpp
@@ -346,6 +346,10 @@ size_t TfLiteMotionPredictorModel::inputLength() const {
     return getTensorBuffer<const float>(mInputR).size();
 }
 
+size_t TfLiteMotionPredictorModel::outputLength() const {
+    return getTensorBuffer<const float>(mOutputR).size();
+}
+
 std::span<float> TfLiteMotionPredictorModel::inputR() {
     return getTensorBuffer<float>(mInputR);
 }
diff --git a/libs/input/tests/TfLiteMotionPredictor_test.cpp b/libs/input/tests/TfLiteMotionPredictor_test.cpp
index 6e76ac1e52..b5ed9e4430 100644
--- a/libs/input/tests/TfLiteMotionPredictor_test.cpp
+++ b/libs/input/tests/TfLiteMotionPredictor_test.cpp
@@ -139,8 +139,10 @@ TEST(TfLiteMotionPredictorTest, ModelInputOutputLength) {
 
     ASSERT_TRUE(model->invoke());
 
-    ASSERT_EQ(model->outputR().size(), model->outputPhi().size());
-    ASSERT_EQ(model->outputR().size(), model->outputPressure().size());
+    const int outputLength = model->outputLength();
+    ASSERT_EQ(outputLength, model->outputR().size());
+    ASSERT_EQ(outputLength, model->outputPhi().size());
+    ASSERT_EQ(outputLength, model->outputPressure().size());
 }
 
 TEST(TfLiteMotionPredictorTest, ModelOutput) {