added back checks in linear_regression_main.cpp and linear_regression…

…_predict_maain.cpp
mlpack · Feb 28, 2024 · 40676a7 · 40676a7
1 parent d534ede
commit 40676a7
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 1 deletion.
diff --git a/src/mlpack/methods/linear_regression/linear_regression_main.cpp b/src/mlpack/methods/linear_regression/linear_regression_main.cpp
@@ -195,7 +195,18 @@ void BINDING_FUNCTION(util::Params& params, util::Timers& timer)
     timer.Stop("load_test_points");
 
     mat points = std::move(params.Get<mat>("test"));
-
+
+    // Ensure that test file data has the right number of features.
+    try
+    {
+      util::CheckSameDimensionality(points, lr->Parameters().n_elem - 1,
+        "Linear Regression Prediction", "test points");
+    }
+    catch (std::invalid_argument& e)
+    {
+      Log::Fatal << e.what() << std::endl;
+    }
+
     // Perform the predictions using our model.
     rowvec predictions;
     timer.Start("prediction");

diff --git a/src/mlpack/methods/linear_regression/linear_regression_predict_main.cpp b/src/mlpack/methods/linear_regression/linear_regression_predict_main.cpp
@@ -64,6 +64,17 @@ void BINDING_FUNCTION(util::Params& params, util::Timers& timer)
 
   mat points = std::move(params.Get<mat>("test"));
 
+  // Ensure that test file data has the right number of features.
+  try
+  {
+    util::CheckSameDimensionality(points, lr->Parameters().n_elem - 1,
+      "Linear Regression Prediction", "test points");
+  }
+  catch (std::invalid_argument& e)
+  {
+    Log::Fatal << e.what() << std::endl;
+  }
+
   // Perform the predictions using our model.
   rowvec predictions;
   timer.Start("prediction");