d5604eaba321c11c1b9616c283262c4ddea55049,rllib/models/torch/modules/relative_multi_head_attention.py,RelativeMultiHeadAttention,forward,#RelativeMultiHeadAttention#Any#Any#,70

Before Change

        Tau = list(memory.shape)[1] if memory is not None else 0
        if memory is not None:
            memory.requires_grad_(False)
            inputs = torch.cat((memory, inputs), dim=1)

        // Apply the Layer-Norm.
        if self._input_layernorm is not None:
            inputs = self._input_layernorm(inputs)

After Change

        score = score / d**0.5

        // causal mask of the same length as the sequence
        mask = sequence_mask(
            torch.arange(Tau + 1, Tau + T + 1),
            dtype=score.dtype).to(score.device)
        mask = mask[None, :, :, None]

        masked_score = score * mask + 1e30 * (mask.float() - 1.)
        wmat = nn.functional.softmax(masked_score, dim=2)

Instances