diff --git a/src/mlpack/methods/ann/layer/multihead_attention.hpp b/src/mlpack/methods/ann/layer/multihead_attention.hpp
index 8299933f5a..cc21bfce37 100644
--- a/src/mlpack/methods/ann/layer/multihead_attention.hpp
+++ b/src/mlpack/methods/ann/layer/multihead_attention.hpp
@@ -50,7 +50,12 @@ namespace mlpack {
  * The input to this layer is expected to be a sequence of embedding vectors.
  * The embedding size is inferred from inputDimensions[0], and the source
  * sequence length is inferred from inputDimensions[1].  If there are more than
- * 2 dimensions, they are flattened into the source sequence length.
+ * 2 dimensions, they are flattened into the source sequence length.  If
+ * selfAttention is true, then query, key, and value are all the same, so the
+ * input data should be of size [embedDim * seqLen, batchSize].  Otherwise,
+ * the input data should be of size
+ * [embedDim * (2 * srcSeqLen + tgtSeqLen), batchSize].  The
+ * output data will always be of size (embedDim * tgtSeqLen, batchSize)
  *
  * @tparam MatType Type of the input/output data (arma::colvec, arma::mat,
  *         arma::sp_mat or arma::cube).
@@ -76,11 +81,15 @@ class MultiheadAttentionType : public Layer<MatType>
    *
    * @param tgtSeqLen Target sequence length.
    * @param numHeads Number of parallel attention heads.
+   * @param attnMask Two dimensional Attention Mask.  Takes the values [-Inf, 0]
+   * @param keyPaddingMask Key Padding Mask.  Takes the values [-Inf, 0]
    * @param selfAttention Use self-attention; source key, query, and value all
    *     come from the same inputs
    */
   MultiheadAttentionType(const size_t tgtSeqLen,
                          const size_t numHeads,
+                         const MatType& attnMask = MatType(),
+                         const MatType& keyPaddingMask = MatType(),
                          const bool selfAttention = false);
 
   //! Clone the MultiheadAttentionType object. This handles polymorphism
@@ -139,6 +148,11 @@ class MultiheadAttentionType : public Layer<MatType>
   template<typename Archive>
   void serialize(Archive& ar, const uint32_t /* version */);
 
+  //! Get the parameters.
+  MatType const& Parameters() const override { return weights; }
+  //! Modify the parameters.
+  MatType& Parameters() override { return weights; }
+
   //! Get the target sequence length.
   size_t TgtSeqLen() const { return tgtSeqLen; }
   //! Modify the target sequence length.
@@ -183,9 +197,29 @@ class MultiheadAttentionType : public Layer<MatType>
           << " [EmbeddingDim,SequenceLen]" << std::endl;
     }
     embedDim = this->inputDimensions[0];
-    srcSeqLen = this->inputDimensions[1];
-    for (size_t i=2; i<this->inputDimensions.size(); i++) {
-      srcSeqLen *= this->inputDimensions[i];
+    if (selfAttention) {
+      // for selfAttention, query, key, and value
+      // all come from the same values
+      // also, srcSeqLen must == tgtSeqLen
+      srcSeqLen = this->inputDimensions[1];
+      for (size_t i=2; i<this->inputDimensions.size(); i++) {
+        srcSeqLen *= this->inputDimensions[i];
+      }
+      if (srcSeqLen != tgtSeqLen) {
+        Log::Fatal << "If using selfAttention, srcSeqLen must equal tgtSeqLen"
+                   << std::endl;
+      }
+    }
+    else {
+      // if we are not using self attention (where query == key == value),
+      // then the source sequence needs have
+      // query = tgtSeqLen
+      // key = srcSeqLen
+      // value = srcSeqLen
+      if ((this->inputDimensions[1] - tgtSeqLen) % 2 != 0) {
+        Log::Fatal << "input dimension 1 is invalid." << std::endl;
+      }
+      srcSeqLen = (this->inputDimensions[1] - tgtSeqLen) / 2;
     }
     if (embedDim % numHeads != 0)
     {
@@ -197,9 +231,17 @@ class MultiheadAttentionType : public Layer<MatType>
 
     // This returns the output as a 2-dimensional (embedDim * tgtSeqLen)
     // matrix.
-    this->outputDimensions = std::vector<size_t>(2, 1);
+    this->outputDimensions = std::vector<size_t>(this->inputDimensions.size(), 1);
     this->outputDimensions[0] = embedDim;
     this->outputDimensions[1] = tgtSeqLen;
+    for (size_t i=2; i<this->outputDimensions.size(); i++) {
+      this->outputDimensions[i] = this->inputDimensions[i];
+      if (this->outputDimensions[1] % this->inputDimensions[i] != 0) {
+        Log::Fatal << "tgtSeqLen " << tgtSeqLen << " not divisible by extra "
+            << "dimension " << this->inputDimensions[i] << std::endl;
+      }
+      this->outputDimensions[1] /= this->inputDimensions[i];
+    }
   }
 
   size_t InputShape() const
@@ -226,10 +268,11 @@ class MultiheadAttentionType : public Layer<MatType>
   //! Dimensionality of each head.
   size_t headDim;
 
-  //! Two dimensional Attention Mask of shape (tgtSeqLen, srcSeqLen).
+  //! Two dimensional Attention Mask of shape (tgtSeqLen, srcSeqLen).  Takes
+  //! the values [-Inf, 0]
   MatType attnMask;
 
-  //! Key Padding Mask.
+  //! Key Padding Mask.  Takes the values [-Inf, 0]
   MatType keyPaddingMask;
 
   //! Whether or not self-attention is used (source key, value, and query all
diff --git a/src/mlpack/methods/ann/layer/multihead_attention_impl.hpp b/src/mlpack/methods/ann/layer/multihead_attention_impl.hpp
index 61e1e25b3c..d804a7d5cd 100644
--- a/src/mlpack/methods/ann/layer/multihead_attention_impl.hpp
+++ b/src/mlpack/methods/ann/layer/multihead_attention_impl.hpp
@@ -38,11 +38,15 @@ MultiheadAttentionType<MatType, RegularizerType>::
 MultiheadAttentionType(
     const size_t tgtSeqLen,
     const size_t numHeads,
+    const MatType& attnmask,
+    const MatType& keypaddingmask,
     const bool selfAttention) :
     tgtSeqLen(tgtSeqLen),
     srcSeqLen(0),
     embedDim(0),
     numHeads(numHeads),
+    attnMask(attnmask),
+    keyPaddingMask(keypaddingmask),
     selfAttention(selfAttention)
 {
 }
@@ -141,40 +145,24 @@ Forward(const MatType& input, MatType& output)
 
   // Apply the attention mask if provided. The attention mask is used to black-
   // out future sequences and generally used in Encoder-Decoder attention.
-  // The attention mask has elements 0 or 1.
+  // The attention mask has elements -inf or 0.
   // The shape of the attention mask : (tgtSeqLen, srcSeqLen).
   if (!attnMask.is_empty())
   {
     if (attnMask.n_rows != tgtSeqLen || attnMask.n_cols != srcSeqLen)
       Log::Fatal << "The size of the 'attn_mask' is not correct.\n";
-    // not sure if there is a better way to do this.  now that the mask
-    // is 0 or 1, we can't simply add, but converting first seems slow
-    const arma::uword rows = tgtSeqLen;
-    const arma::uword cols = srcSeqLen;
-    for (arma::uword i=0; i<rows; i++) {
-      for (arma::uword j=0; j<cols; j++) {
-        if (attnMask.at(i, j) == 0) {
-          scores.tube(i, j).fill(-INFINITY);
-        }
-      }
-    }
+    scores.each_slice() += attnMask;
   }
 
   // Apply the key padding mask when provided. It blacks-out any particular
   // word in the sequence.
-  // The key padding mask has elements 0 or 1.
+  // The key padding mask has elements -inf or 0
   // The shape of keyPaddingMask : (1, srcSeqLen).
   if (!keyPaddingMask.is_empty())
   {
     if (keyPaddingMask.n_rows != 1 || keyPaddingMask.n_cols != srcSeqLen)
         Log::Fatal << "The size of the 'keyPaddingMask' is not correct.\n";
-    // not sure if there is a better way to do this.  now that the mask
-    // is 0 or 1, we can't simply add, but converting first seems slow
-    for (arma::uword pos=0; pos<keyPaddingMask.n_elem; pos++) {
-      if (keyPaddingMask.at(pos) == 0) {
-        scores.tube(0, pos, tgtSeqLen, pos).fill(-INFINITY);
-      }
-    }
+    scores.each_slice() += arma::repmat(keyPaddingMask, tgtSeqLen, 1);
   }
 
   for (size_t i = 0; i < numHeads * batchSize; ++i)
@@ -491,11 +479,11 @@ serialize(Archive& ar, const uint32_t /* version */)
   ar(CEREAL_NVP(selfAttention));
   ar(CEREAL_NVP(softmax));
   ar(CEREAL_NVP(regularizer));
+  ar(CEREAL_NVP(attnMask));
+  ar(CEREAL_NVP(keyPaddingMask));
 
   if (Archive::is_loading::value)
   {
-    attnMask.clear();
-    keyPaddingMask.clear();
     queryWt.clear();
     keyWt.clear();
     valueWt.clear();
diff --git a/src/mlpack/tests/ann/layer/multihead_attention.cpp b/src/mlpack/tests/ann/layer/multihead_attention.cpp
new file mode 100644
index 0000000000..154d2f11d6
--- /dev/null
+++ b/src/mlpack/tests/ann/layer/multihead_attention.cpp
@@ -0,0 +1,237 @@
+/**
+ * @file tests/layer/multihead_attention.cpp
+ * @author Marcus Edel
+ * @author Praveen Ch
+ *
+ * Tests the multihead_attention layer.
+ *
+ * mlpack is free software; you may redistribute it and/or modify it under the
+ * terms of the 3-clause BSD license.  You should have received a copy of the
+ * 3-clause BSD license along with mlpack.  If not, see
+ * http://www.opensource.org/licenses/BSD-3-Clause for more information.
+ */
+#include <mlpack/core.hpp>
+#include <mlpack/methods/ann.hpp>
+
+#include "../../test_catch_tools.hpp"
+#include "../../catch.hpp"
+#include "../../serialization.hpp"
+#include "../ann_test_tools.hpp"
+
+using namespace mlpack;
+
+/**
+ * Simple Multihead Attention test.
+ */
+TEST_CASE("SimpleMultiheadAttentionTest", "[ANNLayerTest]")
+{
+  size_t tLen = 5;
+  size_t sLen = tLen;
+  size_t embedDim = 4;
+  size_t numHeads = 2;
+  size_t bsz = 3;
+
+  arma::mat query = 0.1 * arma::randu(embedDim * tLen, bsz);
+  arma::mat output;
+
+  arma::mat attnMask = arma::zeros(tLen, sLen);
+  for (size_t i = 0; i < tLen; ++i)
+  {
+    for (size_t j = 0; j < sLen; ++j)
+    {
+      if (i < j)
+        attnMask(i, j) = std::numeric_limits<double>::lowest();
+    }
+  }
+
+  arma::mat keyPaddingMask = arma::zeros(1, sLen);
+  keyPaddingMask(sLen - 1) = std::numeric_limits<double>::lowest();
+
+  MultiheadAttention module(tLen, numHeads);
+  module.InputDimensions() = std::vector<size_t>({ embedDim, 2 * sLen + tLen });
+  module.ComputeOutputDimensions();
+  arma::mat weights(module.WeightSize(), 1);
+  weights.randu();
+  module.SetWeights(weights.memptr());
+
+  module.AttentionMask() = attnMask;
+  module.KeyPaddingMask() = keyPaddingMask;
+
+  // Forward test.
+  arma::mat input = arma::join_cols(arma::join_cols(query, query), query);
+
+  module.Forward(input, output);
+  REQUIRE(output.n_rows == embedDim * tLen);
+  REQUIRE(output.n_cols == bsz);
+
+  // Backward test.
+  arma::mat gy = 0.01 * arma::randu(embedDim * tLen, bsz);
+  arma::mat g;
+  module.Backward(input, output, gy, g);
+  REQUIRE(g.n_rows == input.n_rows);
+  REQUIRE(g.n_cols == input.n_cols);
+
+  // Gradient test.
+  arma::mat error = 0.05 * arma::randu(embedDim * tLen, bsz);
+  arma::mat gradient;
+  module.Gradient(input, error, gradient);
+  REQUIRE(gradient.n_rows == module.Parameters().n_rows);
+  REQUIRE(gradient.n_cols == module.Parameters().n_cols);
+}
+
+/**
+ * Jacobian MultiheadAttention module test.
+ */
+TEST_CASE("JacobianMultiheadAttentionTest", "[ANNLayerTest]")
+{
+  // Check when query = key = value.
+  for (size_t i = 0; i < 5; ++i)
+  {
+    const size_t tgtSeqLen = 2;
+    const size_t embedDim = 4;
+    const size_t nHeads = 2;
+    const size_t batchSize = 1;
+
+    arma::mat query = arma::randu(embedDim * tgtSeqLen, batchSize);
+    arma::mat input = arma::join_cols(arma::join_cols(query, query), query);
+
+    MultiheadAttention module(tgtSeqLen, nHeads);
+    module.InputDimensions() = std::vector<size_t>({ embedDim, 3 * tgtSeqLen });
+    module.ComputeOutputDimensions();
+    arma::mat weights(module.WeightSize(), 1);
+    weights.randu();
+    module.SetWeights(weights.memptr());
+
+    double error = CustomJacobianTest(module, input);
+    REQUIRE(error <= 1e-5);
+  }
+
+  // Check when key = value.
+  for (size_t i = 0; i < 5; ++i)
+  {
+    const size_t tgtSeqLen = 2;
+    const size_t srcSeqLen = RandInt(2, 5);
+    const size_t embedDim = 4;
+    const size_t nHeads = 2;
+    const size_t batchSize = 1;
+
+    arma::mat query = arma::randu(embedDim * tgtSeqLen, batchSize);
+    arma::mat key = 0.091 * arma::randu(embedDim * srcSeqLen, batchSize);
+    arma::mat input = arma::join_cols(arma::join_cols(query, key), key);
+
+    MultiheadAttention module(tgtSeqLen, nHeads);
+    module.InputDimensions() = std::vector<size_t>({ embedDim, 2 * srcSeqLen + tgtSeqLen });
+    module.ComputeOutputDimensions();
+    arma::mat weights(module.WeightSize(), 1);
+    weights.randu();
+    module.SetWeights(weights.memptr());
+
+    double error = CustomJacobianTest(module, input);
+    REQUIRE(error <= 1e-5);
+  }
+
+  // Check when query, key and value are not same.
+  for (size_t i = 0; i < 5; ++i)
+  {
+    const size_t tgtSeqLen = 2;
+    const size_t srcSeqLen = RandInt(2, 5);
+    const size_t embedDim = 4;
+    const size_t nHeads = 2;
+    const size_t batchSize = 1;
+
+    arma::mat query = arma::randu(embedDim * tgtSeqLen, batchSize);
+    arma::mat key = 0.091 * arma::randu(embedDim * srcSeqLen, batchSize);
+    arma::mat value = 0.045 * arma::randu(embedDim * srcSeqLen, batchSize);
+    arma::mat input = arma::join_cols(arma::join_cols(query, key), value);
+
+    MultiheadAttention module(tgtSeqLen, nHeads);
+    module.InputDimensions() = std::vector<size_t>({ embedDim, 2 * srcSeqLen + tgtSeqLen });
+    module.ComputeOutputDimensions();
+    arma::mat weights(module.WeightSize(), 1);
+    weights.randu();
+    module.SetWeights(weights.memptr());
+
+    double error = JacobianTest(module, input);
+    REQUIRE(error <= 1e-5);
+  }
+}
+
+
+/**
+ * Numerical gradient test for MultiheadAttention layer.
+ */
+TEST_CASE("GradientMultiheadAttentionTest", "[ANNLayerTest]")
+{
+  struct GradientFunction
+  {
+    GradientFunction() :
+        tgtSeqLen(2),
+        srcSeqLen(2),
+        embedDim(4),
+        nHeads(2),
+        vocabSize(5),
+        batchSize(2)
+    {
+      input = arma::randu(embedDim * (tgtSeqLen + 2 * srcSeqLen), batchSize);
+      target = arma::zeros(vocabSize, batchSize);
+      for (size_t i = 0; i < target.n_elem; ++i)
+      {
+        const size_t label = RandInt(1, vocabSize);
+        target(i) = label;
+      }
+
+      attnMask = arma::zeros(tgtSeqLen, srcSeqLen);
+      for (size_t i = 0; i < tgtSeqLen; ++i)
+      {
+        for (size_t j = 0; j < srcSeqLen; ++j)
+        {
+          if (i < j)
+            attnMask(i, j) = std::numeric_limits<double>::lowest();
+        }
+      }
+
+      keyPaddingMask = arma::zeros(1, srcSeqLen);
+      keyPaddingMask(srcSeqLen - 1) = std::numeric_limits<double>::lowest();
+
+      model = new FFN<NegativeLogLikelihood, XavierInitialization>();
+      model->InputDimensions() = {embedDim, srcSeqLen * 2 + tgtSeqLen};
+      model->ResetData(input, target);
+      // attnModule = new MultiheadAttention(tgtSeqLen, srcSeqLen, embedDim,
+      //     nHeads);
+      // attnModule->AttentionMask() = attnMask;
+      // attnModule->KeyPaddingMask() = keyPaddingMask;
+      // model->Add(attnModule);
+      model->Add<MultiheadAttention>(tgtSeqLen, nHeads,
+                                     attnMask, keyPaddingMask);
+      model->Add<Linear>(vocabSize);
+      model->Add<LogSoftMax>();
+    }
+
+    ~GradientFunction()
+    {
+      delete model;
+    }
+
+    double Gradient(arma::mat& gradient) const
+    {
+      double error = model->Evaluate(model->Parameters(), 0, batchSize);
+      model->Gradient(model->Parameters(), 0, gradient, batchSize);
+      return error;
+    }
+
+    arma::mat& Parameters() { return model->Parameters(); }
+
+    FFN<NegativeLogLikelihood, XavierInitialization>* model;
+    // MultiheadAttention* attnModule;
+
+    arma::mat input, target, attnMask, keyPaddingMask;
+    const size_t tgtSeqLen;
+    const size_t srcSeqLen;
+    const size_t embedDim;
+    const size_t nHeads;
+    const size_t vocabSize;
+    const size_t batchSize;
+  } function;
+
+  REQUIRE(CheckGradient(function) <= 3e-06);
+}
diff --git a/src/mlpack/tests/ann/layer_test.cpp b/src/mlpack/tests/ann/layer_test.cpp
index 320fb451fc..2a12528566 100644
--- a/src/mlpack/tests/ann/layer_test.cpp
+++ b/src/mlpack/tests/ann/layer_test.cpp
@@ -42,3 +42,4 @@
 #include "layer/softmin.cpp"
 #include "layer/ftswish.cpp"
 #include "layer/layer_norm.cpp"
+#include "layer/multihead_attention.cpp"
diff --git a/src/mlpack/tests/ann/not_adapted/ann_layer_test.cpp b/src/mlpack/tests/ann/not_adapted/ann_layer_test.cpp
index 4d37a61ecb..e376586190 100644
--- a/src/mlpack/tests/ann/not_adapted/ann_layer_test.cpp
+++ b/src/mlpack/tests/ann/not_adapted/ann_layer_test.cpp
@@ -3980,204 +3980,6 @@ TEST_CASE("JacobianPositionalEncodingTest", "[ANNLayerTest]")
 }
 */
 
-/**
- * Simple Multihead Attention test.
- */
-TEST_CASE("SimpleMultiheadAttentionTest", "[ANNLayerTest]")
-{
-  size_t tLen = 5;
-  size_t sLen = tLen;
-  size_t embedDim = 4;
-  size_t numHeads = 2;
-  size_t bsz = 3;
-
-  arma::mat query = 0.1 * arma::randu(embedDim * tLen, bsz);
-  arma::mat output;
-
-  arma::mat attnMask = arma::zeros(tLen, sLen);
-  for (size_t i = 0; i < tLen; ++i)
-  {
-    for (size_t j = 0; j < sLen; ++j)
-    {
-      if (i < j)
-        attnMask(i, j) = std::numeric_limits<double>::lowest();
-    }
-  }
-
-  arma::mat keyPaddingMask = arma::zeros(1, sLen);
-  keyPaddingMask(sLen - 1) = std::numeric_limits<double>::lowest();
-
-  MultiheadAttention module(tLen, sLen, embedDim, numHeads);
-  module.AttentionMask() = attnMask;
-  module.KeyPaddingMask() = keyPaddingMask;
-  module.Reset();
-  module.Parameters().randu();
-
-  // Forward test.
-  arma::mat input = arma::join_cols(arma::join_cols(query, query), query);
-
-  module.Forward(input, output);
-  REQUIRE(output.n_rows == embedDim * tLen);
-  REQUIRE(output.n_cols == bsz);
-
-  // Backward test.
-  arma::mat gy = 0.01 * arma::randu(embedDim * tLen, bsz);
-  arma::mat g;
-  module.Backward(input, gy, g);
-  REQUIRE(g.n_rows == input.n_rows);
-  REQUIRE(g.n_cols == input.n_cols);
-
-  // Gradient test.
-  arma::mat error = 0.05 * arma::randu(embedDim * tLen, bsz);
-  arma::mat gradient;
-  module.Gradient(input, error, gradient);
-  REQUIRE(gradient.n_rows == module.Parameters().n_rows);
-  REQUIRE(gradient.n_cols == module.Parameters().n_cols);
-}
-
-/**
- * Jacobian MultiheadAttention module test.
- */
-TEST_CASE("JacobianMultiheadAttentionTest", "[ANNLayerTest]")
-{
-  // Check when query = key = value.
-  for (size_t i = 0; i < 5; ++i)
-  {
-    const size_t tgtSeqLen = 2;
-    const size_t embedDim = 4;
-    const size_t nHeads = 2;
-    const size_t batchSize = 1;
-
-    arma::mat query = arma::randu(embedDim * tgtSeqLen, batchSize);
-    arma::mat input = arma::join_cols(arma::join_cols(query, query), query);
-
-    MultiheadAttention module(tgtSeqLen, tgtSeqLen, embedDim, nHeads);
-    module.Parameters().randu();
-
-    double error = CustomJacobianTest(module, input);
-    REQUIRE(error <= 1e-5);
-  }
-
-  // Check when key = value.
-  for (size_t i = 0; i < 5; ++i)
-  {
-    const size_t tgtSeqLen = 2;
-    const size_t srcSeqLen = RandInt(2, 5);
-    const size_t embedDim = 4;
-    const size_t nHeads = 2;
-    const size_t batchSize = 1;
-
-    arma::mat query = arma::randu(embedDim * tgtSeqLen, batchSize);
-    arma::mat key = 0.091 * arma::randu(embedDim * srcSeqLen, batchSize);
-    arma::mat input = arma::join_cols(arma::join_cols(query, key), key);
-
-    MultiheadAttention module(tgtSeqLen, srcSeqLen, embedDim, nHeads);
-    module.Parameters().randu();
-
-    double error = CustomJacobianTest(module, input);
-    REQUIRE(error <= 1e-5);
-  }
-
-  // Check when query, key and value are not same.
-  for (size_t i = 0; i < 5; ++i)
-  {
-    const size_t tgtSeqLen = 2;
-    const size_t srcSeqLen = RandInt(2, 5);
-    const size_t embedDim = 4;
-    const size_t nHeads = 2;
-    const size_t batchSize = 1;
-
-    arma::mat query = arma::randu(embedDim * tgtSeqLen, batchSize);
-    arma::mat key = 0.091 * arma::randu(embedDim * srcSeqLen, batchSize);
-    arma::mat value = 0.045 * arma::randu(embedDim * srcSeqLen, batchSize);
-    arma::mat input = arma::join_cols(arma::join_cols(query, key), value);
-
-    MultiheadAttention module(tgtSeqLen, srcSeqLen, embedDim, nHeads);
-    module.Parameters().randu();
-
-    double error = JacobianTest(module, input);
-    REQUIRE(error <= 1e-5);
-  }
-}
-
-
-/**
- * Numerical gradient test for MultiheadAttention layer.
- */
-TEST_CASE("GradientMultiheadAttentionTest", "[ANNLayerTest]")
-{
-  struct GradientFunction
-  {
-    GradientFunction() :
-        tgtSeqLen(2),
-        srcSeqLen(2),
-        embedDim(4),
-        nHeads(2),
-        vocabSize(5),
-        batchSize(2)
-    {
-      input = arma::randu(embedDim * (tgtSeqLen + 2 * srcSeqLen), batchSize);
-      target = arma::zeros(vocabSize, batchSize);
-      for (size_t i = 0; i < target.n_elem; ++i)
-      {
-        const size_t label = RandInt(1, vocabSize);
-        target(i) = label;
-      }
-
-      attnMask = arma::zeros(tgtSeqLen, srcSeqLen);
-      for (size_t i = 0; i < tgtSeqLen; ++i)
-      {
-        for (size_t j = 0; j < srcSeqLen; ++j)
-        {
-          if (i < j)
-            attnMask(i, j) = std::numeric_limits<double>::lowest();
-        }
-      }
-
-      keyPaddingMask = arma::zeros(1, srcSeqLen);
-      keyPaddingMask(srcSeqLen - 1) = std::numeric_limits<double>::lowest();
-
-      model = new FFN<NegativeLogLikelihood, XavierInitialization>();
-      model->ResetData(input, target);
-      // attnModule = new MultiheadAttention(tgtSeqLen, srcSeqLen, embedDim,
-      //     nHeads);
-      // attnModule->AttentionMask() = attnMask;
-      // attnModule->KeyPaddingMask() = keyPaddingMask;
-      // model->Add(attnModule);
-      model->Add<MultiheadAttention>(tgtSeqLen, srcSeqLen, embedDim, nHeads,
-          attnMask, keyPaddingMask);
-      model->Add<Linear>(embedDim * tgtSeqLen, vocabSize);
-      model->Add<LogSoftMax>();
-    }
-
-    ~GradientFunction()
-    {
-      delete model;
-    }
-
-    double Gradient(arma::mat& gradient) const
-    {
-      double error = model->Evaluate(model->Parameters(), 0, batchSize);
-      model->Gradient(model->Parameters(), 0, gradient, batchSize);
-      return error;
-    }
-
-    arma::mat& Parameters() { return model->Parameters(); }
-
-    FFN<NegativeLogLikelihood, XavierInitialization>* model;
-    // MultiheadAttention* attnModule;
-
-    arma::mat input, target, attnMask, keyPaddingMask;
-    const size_t tgtSeqLen;
-    const size_t srcSeqLen;
-    const size_t embedDim;
-    const size_t nHeads;
-    const size_t vocabSize;
-    const size_t batchSize;
-  } function;
-
-  REQUIRE(CheckGradient(function) <= 3e-06);
-}
 
 /**
  * Simple tests for instance normalization layer.