Fix output_layer call
@@ -102,7 +102,7 @@ class SelfAttentionDecoder(Decoder):
memory=cache["memory"],
memory_sequence_length=None)
outputs = outputs[:, -1:, :]
- logits = output_layer(outputs, vocab_size)
+ logits = output_layer(outputs)
return logits, cache
return _impl
