Fix numClasses member in SoftmaxRegression::Train()

HISTORY.md

@@ -8,6 +8,9 @@
 
   * More robust detection of C++17 mode in the MSVC "compiler" (#3555).
 
+  * Fix setting number of classes correctly in `SoftmaxRegression::Train()`
+    (#3553).
+
 ### mlpack 4.2.1
 ###### 2023-09-05
   * Reinforcement Learning: Gaussian noise (#3515).

src/mlpack/methods/softmax_regression/softmax_regression_impl.hpp

@@ -190,6 +190,7 @@ double SoftmaxRegression::Train(const arma::mat& data,
 
   // Train the model.
   const double out = optimizer.Optimize(regressor, parameters);
+  this->numClasses = numClasses;
 
   Log::Info << "SoftmaxRegression::SoftmaxRegression(): final objective of "
             << "trained model is " << out << "." << std::endl;
@@ -211,6 +212,7 @@ double SoftmaxRegression::Train(const arma::mat& data,
 
   // Train the model.
   const double out = optimizer.Optimize(regressor, parameters, callbacks...);
+  this->numClasses = numClasses;
 
   Log::Info << "SoftmaxRegression::SoftmaxRegression(): final objective of "
             << "trained model is " << out << "." << std::endl;

src/mlpack/tests/softmax_regression_test.cpp

@@ -672,3 +672,32 @@ TEST_CASE("SoftmaxRegressionComputeProbabilitiesAndLabelsTest",
     REQUIRE(testLabels(i) == labels(i));
   }
 }
+
+TEST_CASE("SoftmaxImmediateTrainTest", "[SoftmaxRegressionTest]")
+{
+  // Initialize a random dataset.
+  const size_t numClasses = 3;
+  const size_t inputSize = 10;
+  const size_t points = 500;
+  arma::mat data;
+  data.randu(inputSize, points);
+
+  // Create random class labels.
+  arma::Row<size_t> labels(points);
+  for (size_t i = 0; i < points; ++i)
+    labels(i) = RandInt(0, numClasses);
+
+  // Train without setting any parameters to the constructor.
+  SoftmaxRegression sr;
+  sr.Train(data, labels, numClasses);
+
+  // Now classify some points.
+  // This just makes sure that the model can successfully make predictions at
+  // all (i.e. no exception thrown).
+  arma::Row<size_t> predictions;
+  sr.Classify(data, predictions);
+
+  REQUIRE(predictions.n_elem == labels.n_elem);
+  REQUIRE(arma::all(predictions >= 0));
+  REQUIRE(arma::all(predictions <= 2));
+}