    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.permute(0, 2, 1)
        out = self.conv(input).permute(0, 2, 1)
        out = self.activatoin(self.norm(out))
        return out
forward for training
modules/wenet_extractor/transducer/predictor.py