    def forward_step(
        self,
        input: torch.Tensor,
        padding: torch.Tensor,
        cache: List[torch.Tensor],
    ) -> Tuple[torch.Tensor, List[torch.Tensor]]:
        """forward step for inference
        Args:
            input (torch.Tensor): [batch_size, time_step=1]
            padding (torch.Tensor): [batch_size,1], 1 is padding value
            cache: for embedding predictor, cache[0] == history
        """
        assert input.size(1) == 1
        assert len(cache) == 1
        history = cache[0]
        assert history.size(1) == self.context_size - 1
        input = self.embed(input)  # [bs, 1, embed]
        input = self.embed_dropout(input)
        context_input = torch.cat((history, input), dim=1)
        input_expand = context_input.unsqueeze(1).unsqueeze(
            2
        )  # [bs, 1, 1, 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
        )

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

        output = self.ffn(output)
        output = self.norm(output)
        output = self.activatoin(output)
        new_cache = context_input[:, 1:, :]
        # TODO(Mddct): we need padding new_cache in future
        # new_cache = ApplyPadding(history, padding, new_cache)
        return (output, [new_cache])
forward step for inference
Args:
input (torch.Tensor): [batch_size, time_step=1]
padding (torch.Tensor): [batch_size,1], 1 is padding value
cache: for embedding predictor, cache[0] == history
modules/wenet_extractor/transducer/predictor.py