# Buggy code
def forward(
    self,
    prev_output_tokens,
    encoder_out: Optional[EncoderOut] = None,
    incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]] = None,
    features_only: bool = False,
    alignment_layer: Optional[int] = None,
    alignment_heads: Optional[int] = None,
    src_lengths: Optional[Any] = None,
    return_all_hiddens: bool = False,
):
    """
    Args:
        prev_output_tokens (LongTensor): previous decoder outputs of shape
            `(batch, tgt_len)`, for teacher forcing
        encoder_out (optional): output from the encoder, used for
            encoder-side attention
        incremental_state (dict): dictionary used for storing state during
            :ref:`Incremental decoding`
        features_only (bool, optional): only return features without
            applying output layer (default: False).

    Returns:
        tuple:
            - the decoder's output of shape `(batch, tgt_len, vocab)`
            - a dictionary with any model-specific outputs
    """
    x, extra = self.extract_features(
        prev_output_tokens,
        encoder_out=encoder_out,
        incremental_state=incremental_state,
        alignment_layer=alignment_layer,
        alignment_heads=alignment_heads,
    )
    if not features_only:
        x = self.output_layer(x)
    return x, extra


# Traceback (TypeError)
Traceback (most recent call last):
File ".../bin/fairseq-train", line 11, in <module>
load_entry_point('fairseq', 'console_scripts', 'fairseq-train')()
File ".../fairseq_cli/train.py", line 351, in cli_main
distributed_utils.call_main(args, main)
File ".../fairseq/distributed_utils.py", line 254, in call_main
main(args, **kwargs)
File ".../fairseq_cli/train.py", line 125, in main
valid_losses, should_stop = train(args, trainer, task, epoch_itr)
File "/usr/lib64/python3.6/contextlib.py", line 52, in inner
return func(*args, **kwds)
File ".../fairseq_cli/train.py", line 207, in train
log_output = trainer.train_step(samples)
File "/usr/lib64/python3.6/contextlib.py", line 52, in inner
return func(*args, **kwds)
File ".../fairseq/trainer.py", line 479, in train_step
ignore_grad=is_dummy_batch,
File ".../fairseq/tasks/fairseq_task.py", line 408, in train_step
loss, sample_size, logging_output = criterion(model, sample)
File ".../lib64/python3.6/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
result = self.forward(*input, **kwargs)
File ".../fairseq/criterions/label_smoothed_cross_entropy_with_alignment.py", line 36, in forward
net_output = model(**sample['net_input'])
File ".../lib64/python3.6/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
result = self.forward(*input, **kwargs)
File ".../fairseq/models/transformer_align.py", line 51, in forward
return self.forward_decoder(prev_output_tokens, encoder_out)
File ".../fairseq/models/transformer_align.py", line 75, in forward_decoder
**extra_args,
File ".../lib64/python3.6/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
result = self.forward(*input, **kwargs)
TypeError: forward() got an unexpected keyword argument 'full_context_alignment'

# Fixed code
def forward(
    self,
    prev_output_tokens,
    encoder_out: Optional[EncoderOut] = None,
    incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]] = None,
    features_only: bool = False,
    full_context_alignment: bool = False,
    alignment_layer: Optional[int] = None,
    alignment_heads: Optional[int] = None,
    src_lengths: Optional[Any] = None,
    return_all_hiddens: bool = False,
):
    """
    Args:
        prev_output_tokens (LongTensor): previous decoder outputs of shape
            `(batch, tgt_len)`, for teacher forcing
        encoder_out (optional): output from the encoder, used for
            encoder-side attention
        incremental_state (dict): dictionary used for storing state during
            :ref:`Incremental decoding`
        features_only (bool, optional): only return features without
            applying output layer (default: False).
        full_context_alignment (bool, optional): don't apply
            auto-regressive mask to self-attention (default: False).

    Returns:
        tuple:
            - the decoder's output of shape `(batch, tgt_len, vocab)`
            - a dictionary with any model-specific outputs
    """
    x, extra = self.extract_features(
        prev_output_tokens,
        encoder_out=encoder_out,
        incremental_state=incremental_state,
        full_context_alignment=full_context_alignment,
        alignment_layer=alignment_layer,
        alignment_heads=alignment_heads,
    )
    if not features_only:
        x = self.output_layer(x)
    return x, extra