fix in_channels of LN (#1177)

dmlc · Feb 25, 2020 · bb2db3a · bb2db3a
1 parent a785e6f
commit bb2db3a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/scripts/language_model/model/qa.py b/scripts/language_model/model/qa.py
@@ -43,7 +43,7 @@ def __init__(self, units=768, is_eval=False, prefix=None, params=None):
         with self.name_scope():
             self.dense_0 = nn.Dense(units, activation='tanh', flatten=False)
             self.dense_1 = nn.Dense(1, flatten=False)
-            self.layernorm = nn.LayerNorm(epsilon=1e-12, in_channels=768)
+            self.layernorm = nn.LayerNorm(epsilon=1e-12, in_channels=units)
 
     def __call__(self,
                  hidden_states,