diff --git a/edward2/tensorflow/layers/__init__.py b/edward2/tensorflow/layers/__init__.py
index d9d8e9e5..2a374999 100644
--- a/edward2/tensorflow/layers/__init__.py
+++ b/edward2/tensorflow/layers/__init__.py
@@ -118,6 +118,7 @@
     "MCSigmoidDenseFA",
     "MCSoftmaxDense",
     "MCSoftmaxDenseFA",
+    "MultiHeadMCSoftmaxDenseFA",
     "MixtureLogistic",
     "NCPCategoricalPerturb",
     "NCPNormalOutput",