    def forward(self, input: torch.Tensor, cache: Optional[List[torch.Tensor]] = None):
        """forward for training"""
        input = self.embed(input)  # [bs, seq_len, embed]
        input = self.embed_dropout(input)
        if cache is None:
            zeros = self.init_state(input.size(0), device=input.device)[0]
        else:
            assert len(cache) == 1
            zeros = cache[0]

        input = torch.cat(
            (zeros, input), dim=1
        )  # [bs, context_size-1 + seq_len, embed]

        input = input.unfold(1, self.context_size, 1).permute(
            0, 1, 3, 2
        )  # [bs, seq_len, context_size, embed]
        # multi head pos: [n_head, embed, context_size]
        multi_head_pos = self.pos_embed.weight.view(
            self.num_heads, self.embed_size, self.context_size
        )

        # broadcast dot attenton
        input_expand = input.unsqueeze(2)  # [bs, seq_len, 1, context_size, embed]
        multi_head_pos = multi_head_pos.permute(
            0, 2, 1
        )  # [num_heads, context_size, embed]

        # [bs, seq_len, num_heads, context_size, embed]
        weight = input_expand * multi_head_pos
        weight = weight.sum(dim=-1, keepdim=False).unsqueeze(
            3
        )  # [bs, seq_len, num_heads, 1, context_size]
        output = weight.matmul(input_expand).squeeze(
            dim=3
        )  # [bs, seq_len, num_heads, embed]
        output = output.sum(dim=2)  # [bs, seq_len, embed]
        output = output / (self.num_heads * self.context_size)

        output = self.ffn(output)
        output = self.norm(output)
        output = self.activatoin(output)
        return output
forward for training
modules/wenet_extractor/transducer/predictor.py