def forward(self, feats: Sequence[Sequence[Tensor]]) -> List[Tensor]:
        """Forward the network. The input is multi-stage multi-unit feature
        maps and the output is a list of heatmaps from multiple stages.

        Args:
            feats (Sequence[Sequence[Tensor]]): Feature maps from multiple
                stages and units.

        Returns:
            List[Tensor]: A list of output heatmaps from multiple stages
                and units.
        """
        out = []
        assert len(feats) == self.num_stages, (
            f'The length of feature maps did not match the '
            f'`num_stages` in {self.__class__.__name__}')
        for feat in feats:
            assert len(feat) == self.num_units, (
                f'The length of feature maps did not match the '
                f'`num_units` in {self.__class__.__name__}')
            for f in feat:
                assert f.shape[1] == self.unit_channels, (
                    f'The number of feature map channels did not match the '
                    f'`unit_channels` in {self.__class__.__name__}')

        for i in range(self.num_stages):
            for j in range(self.num_units):
                y = self.predict_layers[i * self.num_units + j](feats[i][j])
                out.append(y)
        return out
Forward the network. The input is multi-stage multi-unit feature
        maps and the output is a list of heatmaps from multiple stages.

        Args:
            feats (Sequence[Sequence[Tensor]]): Feature maps from multiple
                stages and units.

        Returns:
            List[Tensor]: A list of output heatmaps from multiple stages
                and units.
        
mmpose/mmpose/models/heads/heatmap_heads/mspn_head.py