From 1accad6c299367d0a8dde5aa45b159dca6ccf7bf Mon Sep 17 00:00:00 2001
From: Ryan Curtin <ryan@ratml.org>
Date: Wed, 31 Jan 2024 09:45:32 -0500
Subject: [PATCH] Fix bug when loading grayscale: make sure info.Channels() is
 1.

---
 src/mlpack/core/data/load_image_impl.hpp | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/mlpack/core/data/load_image_impl.hpp b/src/mlpack/core/data/load_image_impl.hpp
index b790782188..a86b84eddf 100644
--- a/src/mlpack/core/data/load_image_impl.hpp
+++ b/src/mlpack/core/data/load_image_impl.hpp
@@ -154,14 +154,16 @@ inline bool LoadImage(const std::string& filename,
 
   info.Width() = tempWidth;
   info.Height() = tempHeight;
-  info.Channels() = tempChannels;
+  // Only set the new number of channels if we didn't force grayscale loading.
+  if (info.Channels() != 1)
+    info.Channels() = tempChannels;
 
   // Copy image into armadillo Mat.
   matrix = arma::Mat<unsigned char>(image, info.Width() * info.Height() *
       info.Channels(), 1, true, true);
 
   // Free the image pointer.
-  free(image);
+  stbi_image_free(image);
   return true;
 }