Add onnx_trace argument for learned embeddings
Summary:
Pull Request resolved:
This argument was missing so we cannot export Transformer if we use learned positional embeddings. See also
@@ -20,6 +20,7 @@ class LearnedPositionalEmbedding(nn.Embedding):
def __init__(self, num_embeddings, embedding_dim, padding_idx, left_pad):
super().__init__(num_embeddings, embedding_dim, padding_idx)
self.left_pad = left_pad
+ self.onnx_trace = False
def forward(self, input, incremental_state=None):
"""Input is expected to be of size [bsz x seqlen]."""
@@ -27,7 +28,7 @@ class LearnedPositionalEmbedding(nn.Embedding):
# positions is the same for every token when decoding a single step
positions = input.data.new(1, 1).fill_(self.padding_idx + input.size(1))
else:
- positions = utils.make_positions(input.data, self.padding_idx, self.left_pad)
+ positions = utils.make_positions(input.data, self.padding_idx, self.left_pad, self.onnx_trace)
return super().forward(positions)
def max_positions(self):
