diff --git a/src/mlpack/tests/custom_layer.hpp b/src/mlpack/tests/custom_layer.hpp
index 8d791b1ebd..3a3eaa78ad 100644
--- a/src/mlpack/tests/custom_layer.hpp
+++ b/src/mlpack/tests/custom_layer.hpp
@@ -26,11 +26,9 @@ namespace ann {
  */
 template <
     class ActivationFunction = LogisticFunction,
-    typename InputDataType = arma::mat,
-    typename OutputDataType = arma::mat
+    typename MatType = arma::mat
 >
-using CustomLayer = BaseLayer<
-    ActivationFunction, InputDataType, OutputDataType>;
+using CustomLayer = BaseLayer<ActivationFunction, MatType>;
 
 } // namespace ann
 } // namespace mlpack