From 33af71dbce383b4809509429d40603aa4e69a1ab Mon Sep 17 00:00:00 2001
From: Nippun Sharma <53967069+NippunSharma@users.noreply.github.com>
Date: Sun, 24 Jan 2021 00:35:21 +0530
Subject: [PATCH] Update src/mlpack/core/util/io.cpp

Co-authored-by: Ryan Curtin <ryan@ratml.org>
---
 src/mlpack/core/util/io.cpp | 32 +++++++++++++++++++++-----------
 1 file changed, 21 insertions(+), 11 deletions(-)

diff --git a/src/mlpack/core/util/io.cpp b/src/mlpack/core/util/io.cpp
index 8687a6df12..dba1bbecd8 100644
--- a/src/mlpack/core/util/io.cpp
+++ b/src/mlpack/core/util/io.cpp
@@ -279,26 +279,36 @@ void IO::CheckInputMatrices()
     std::string paramType = itr->second.cppType;
 
     if (paramType == "arma::mat")
-      IO::CheckInputMatrix<arma::Mat<double>>(
-          IO::GetParam<arma::Mat<double>>(paramName), paramName);
+    {
+      IO::CheckInputMatrix(IO::GetParam<arma::mat>(paramName), paramName);
+    }
     else if (paramType == "arma::Mat<size_t>")
-      IO::CheckInputMatrix<arma::Mat<size_t>>(
+    {
+      IO::CheckInputMatrix(
           IO::GetParam<arma::Mat<size_t>>(paramName), paramName);
+    }
     else if (paramType == "arma::colvec")
-      IO::CheckInputMatrix<arma::Col<double>>(
-          IO::GetParam<arma::Col<double>>(paramName), paramName);
+    {
+      IO::CheckInputMatrix(IO::GetParam<arma::colvec>(paramName), paramName);
+    }
     else if (paramType == "arma::Col<size_t>")
-      IO::CheckInputMatrix<arma::Col<size_t>>(
+    {
+      IO::CheckInputMatrix(
           IO::GetParam<arma::Col<size_t>>(paramName), paramName);
+    }
     else if (paramType == "arma::rowvec")
-      IO::CheckInputMatrix<arma::Row<double>>(
-          IO::GetParam<arma::Row<double>>(paramName), paramName);
+    {
+      IO::CheckInputMatrix(IO::GetParam<arma::rowvec>(paramName), paramName);
+    }
     else if (paramType == "arma::Row<size_t>")
-      IO::CheckInputMatrix<arma::Row<size_t>>(
+    {
+      IO::CheckInputMatrix(
           IO::GetParam<arma::Row<size_t>>(paramName), paramName);
+    }
     else if (paramType == "std::tuple<data::DatasetInfo, arma::mat>")
-      IO::CheckInputMatrix<arma::Mat<double>>(
+    {
+      IO::CheckInputMatrix(
         std::get<1>(IO::GetParam<TupleType>(paramName)), paramName);
+    }
   }
 }
-