# Buggy code
def _append_prev_key_padding_mask(
    key_padding_mask: Optional[Tensor],
    prev_key_padding_mask: Optional[Tensor],
    batch_size: int,
    src_len: int,
    static_kv: bool,
) -> Optional[Tensor]:
    # saved key padding masks have shape (bsz, seq_len)
    if prev_key_padding_mask is not None and static_kv:
        new_key_padding_mask = prev_key_padding_mask
    elif prev_key_padding_mask is not None and key_padding_mask is not None:
        new_key_padding_mask = torch.cat(
            [prev_key_padding_mask.float(), key_padding_mask.float()], dim=1
        )
    # During incremental decoding, as the padding token enters and
    # leaves the frame, there will be a time when prev or current
    # is None
    elif prev_key_padding_mask is not None:
        filler = torch.zeros(batch_size, src_len - prev_key_padding_mask.size(1))
        if prev_key_padding_mask.is_cuda:
            filler = filler.cuda()
        new_key_padding_mask = torch.cat(
            [prev_key_padding_mask.float(), filler.float()], dim=1
        )
    elif key_padding_mask is not None:
        filler = torch.zeros(batch_size, src_len - key_padding_mask.size(1))
        if key_padding_mask.is_cuda:
            filler = filler.cuda()
        new_key_padding_mask = torch.cat(
            [filler.float(), key_padding_mask.float()], dim=1
        )
    else:
        new_key_padding_mask = prev_key_padding_mask
    return new_key_padding_mask


# Traceback (RuntimeError)
Using cache found in /home/vlad/.cache/torch/hub/pytorch_fairseq_master
Loading codes from /home/vlad/.cache/torch/pytorch_fairseq/0695ef328ddefcb8cbcfabc3196182f59c0e41e0468b10cc0db2ae9c91881fcc.bb1be17de4233e13870bd7d6065bfdb03fca0a51dd0f5d0b7edf5c188eda71f1/bpecodes ...
Read 30000 codes from the codes file.
Traceback (most recent call last):
File "/home/vlad/Documents/coding/experiments/paper-analyzer-Tretyak_Internship/papers/experiments/scientific_generation/scripts/paraphrase.py", line 46, in <module>
en2de.translate(['hello'])
File "/home/vlad/.cache/torch/hub/pytorch_fairseq_master/fairseq/hub_utils.py", line 126, in translate
return self.sample(sentences, beam, verbose, **kwargs)
File "/home/vlad/.cache/torch/hub/pytorch_fairseq_master/fairseq/hub_utils.py", line 132, in sample
batched_hypos = self.generate(tokenized_sentences, beam, verbose, **kwargs)
File "/home/vlad/.cache/torch/hub/pytorch_fairseq_master/fairseq/hub_utils.py", line 165, in generate
translations = self.task.inference_step(generator, self.models, batch)
File "/home/vlad/.cache/torch/hub/pytorch_fairseq_master/fairseq/tasks/fairseq_task.py", line 351, in inference_step
return generator.generate(models, sample, prefix_tokens=prefix_tokens)
File "/home/vlad/Documents/envs/p3.6/lib/python3.6/site-packages/torch/autograd/grad_mode.py", line 49, in decorate_no_grad
return func(*args, **kwargs)
File "/home/vlad/.cache/torch/hub/pytorch_fairseq_master/fairseq/sequence_generator.py", line 93, in generate
return self._generate(model, sample, **kwargs)
File "/home/vlad/Documents/envs/p3.6/lib/python3.6/site-packages/torch/autograd/grad_mode.py", line 49, in decorate_no_grad
return func(*args, **kwargs)
File "/home/vlad/.cache/torch/hub/pytorch_fairseq_master/fairseq/sequence_generator.py", line 276, in _generate
tokens[:, :step + 1], encoder_outs, temperature=self.temperature,
File "/home/vlad/Documents/envs/p3.6/lib/python3.6/site-packages/torch/autograd/grad_mode.py", line 49, in decorate_no_grad
return func(*args, **kwargs)
File "/home/vlad/.cache/torch/hub/pytorch_fairseq_master/fairseq/sequence_generator.py", line 549, in forward_decoder
temperature=temperature,
File "/home/vlad/.cache/torch/hub/pytorch_fairseq_master/fairseq/sequence_generator.py", line 568, in _decode_one
tokens, encoder_out=encoder_out, incremental_state=self.incremental_states[model],
File "/home/vlad/.cache/torch/hub/pytorch_fairseq_master/fairseq/models/fairseq_model.py", line 274, in forward_decoder
return self.decoder(prev_output_tokens, **kwargs)
File "/home/vlad/Documents/envs/p3.6/lib/python3.6/site-packages/torch/nn/modules/module.py", line 532, in __call__
result = self.forward(*input, **kwargs)
File "/home/vlad/.cache/torch/hub/pytorch_fairseq_master/fairseq/models/transformer.py", line 704, in forward
alignment_heads=alignment_heads,
File "/home/vlad/.cache/torch/hub/pytorch_fairseq_master/fairseq/models/transformer.py", line 807, in extract_features
need_head_weights=bool((idx == alignment_layer)),
File "/home/vlad/Documents/envs/p3.6/lib/python3.6/site-packages/torch/nn/modules/module.py", line 532, in __call__
result = self.forward(*input, **kwargs)
File "/home/vlad/.cache/torch/hub/pytorch_fairseq_master/fairseq/modules/transformer_layer.py", line 297, in forward
need_head_weights=need_head_weights,
File "/home/vlad/Documents/envs/p3.6/lib/python3.6/site-packages/torch/nn/modules/module.py", line 532, in __call__
result = self.forward(*input, **kwargs)
File "/home/vlad/.cache/torch/hub/pytorch_fairseq_master/fairseq/modules/multihead_attention.py", line 325, in forward
key_padding_mask.unsqueeze(1).unsqueeze(2).to(torch.bool), float("-inf")
RuntimeError: arguments are located on different GPUs at /pytorch/aten/src/THC/generic/THCTensorMasked.cu:28

Process finished with exit code 1

# Fixed code
def _append_prev_key_padding_mask(
    key_padding_mask: Optional[Tensor],
    prev_key_padding_mask: Optional[Tensor],
    batch_size: int,
    src_len: int,
    static_kv: bool,
) -> Optional[Tensor]:
    # saved key padding masks have shape (bsz, seq_len)
    if prev_key_padding_mask is not None and static_kv:
        new_key_padding_mask = prev_key_padding_mask
    elif prev_key_padding_mask is not None and key_padding_mask is not None:
        new_key_padding_mask = torch.cat(
            [prev_key_padding_mask.float(), key_padding_mask.float()], dim=1
        )
    # During incremental decoding, as the padding token enters and
    # leaves the frame, there will be a time when prev or current
    # is None
    elif prev_key_padding_mask is not None:
        filler = torch.zeros(
            (batch_size, src_len - prev_key_padding_mask.size(1)),
            device=prev_key_padding_mask.device,
        )
        new_key_padding_mask = torch.cat(
            [prev_key_padding_mask.float(), filler.float()], dim=1
        )
    elif key_padding_mask is not None:
        filler = torch.zeros(
            (batch_size, src_len - key_padding_mask.size(1)),
            device=key_padding_mask.device,
        )
        new_key_padding_mask = torch.cat(
            [filler.float(), key_padding_mask.float()], dim=1
        )
    else:
        new_key_padding_mask = prev_key_padding_mask
    return new_key_padding_mask