diff --git a/configure.ac b/configure.ac
index 07813e45..229f71f6 100644
--- a/configure.ac
+++ b/configure.ac
@@ -10,7 +10,7 @@
 
 # Initialize autoconf.
 AC_PREREQ([2.69])
-AC_INIT([RidgeRun inference library],[0.4.0],[https://github.com/RidgeRun/r2inference/issues],[r2inference])
+AC_INIT([RidgeRun inference library],[0.4.1],[https://github.com/RidgeRun/r2inference/issues],[r2inference])
 
 # Initialize our build utils
 RR_INIT
diff --git a/r2i/tflite/engine.cc b/r2i/tflite/engine.cc
index 0a63b7fd..3bdd99eb 100644
--- a/r2i/tflite/engine.cc
+++ b/r2i/tflite/engine.cc
@@ -201,7 +201,7 @@ std::shared_ptr<r2i::IPrediction> Engine::Predict (std::shared_ptr<r2i::IFrame>
 
   int output = this->interpreter->outputs()[0];
   TfLiteIntArray *output_dims = this->interpreter->tensor(output)->dims;
-  auto output_size = output_dims->data[output_dims->size - 1];
+  auto output_size = output_dims->data[output_dims->size - 1] * sizeof(float);
   auto *tensor_data = this->interpreter->typed_output_tensor<float>(0);
   prediction->SetTensorValues(tensor_data, output_size);