Method forward_layers_checkpointed

OSUM/wenet/LLM/decoder.py:151–157

(self, xs: torch.Tensor,
                                    att_mask: torch.Tensor,
                                    pos_emb: torch.Tensor)

    @torch.jit.ignore(drop=True)
    def forward_layers_checkpointed(self, xs: torch.Tensor,
                                    att_mask: torch.Tensor,
                                    pos_emb: torch.Tensor) -> torch.Tensor:
        for layer in self.decoders:
            xs, _, _, _ = ckpt.checkpoint(layer.__call__, xs, att_mask,
                                          pos_emb)
        return xs

    @property
    def hidden_size(self):