diff --git a/flownmt/nnet/positional_encoding.py b/flownmt/nnet/positional_encoding.py
index 05c3574..adfb618 100644
--- a/flownmt/nnet/positional_encoding.py
+++ b/flownmt/nnet/positional_encoding.py
@@ -46,7 +46,7 @@ def forward(self, x):
             # recompute/expand embeddings if needed
             self.weights = PositionalEncoding.get_embedding(
                 max_pos,
-                self.embedding_dim,
+                self.encoding_dim,
                 self.padding_idx,
             )
         self.weights = self.weights.type_as(self._float_tensor)