    def forward(
        self,
        source: torch.Tensor,
        target_list: Optional[List[torch.Tensor]] = None,
        padding_mask: Optional[torch.Tensor] = None,
        mask: bool = True,
        features_only: bool = False,
        output_layer: Optional[int] = None,
    ) -> Dict[str, torch.Tensor]:
        """output layer is 1-based"""
        features = self.forward_features(source)

        features_pen = features.float().pow(2).mean()

        features = features.transpose(1, 2)
        features = self.layer_norm(features)
        unmasked_features = features.clone()

        if padding_mask is not None:
            padding_mask = self.forward_padding_mask(features, padding_mask)

        if self.post_extract_proj is not None:
            features = self.post_extract_proj(features)

        features = self.dropout_input(features)
        unmasked_features = self.dropout_features(unmasked_features)

        if mask:
            x, mask_indices = self.apply_mask(features, padding_mask, target_list)
        else:
            x = features
            mask_indices = None

        # feature: (B, T, D), float
        # target: (B, T), long
        # x: (B, T, D), float
        # padding_mask: (B, T), bool
        # mask_indices: (B, T), bool

        def align_size_sum(feat1, pad1, feat2):
            assert (
                abs(feat1.size(1) - feat2.size(1)) < 10
            ), "misaligned results for feat1 and feat2 of size {} - {}".format(
                feat1.size(1), feat2.size(1)
            )
            common_size = min(feat1.size(1), feat2.size(1))

            return (
                feat1[:, :common_size] + feat2[:, :common_size],
                pad1[:, :common_size],
            )

        # process encoders
        res_outputs = []  # final output for different resolution
        multi_mask_indices = []  # mask indices for different resolution
        residuals = []  # record the x in encoders
        padding_masks = []  # final padding masks
        # The encoder has (self.label_nums - 1) blocks
        for i in range(self.label_nums - 1):
            x, _ = self.encoders[i](x, padding_mask=padding_mask, layer=None)
            residuals.append(x)
            x, padding_mask, mask_indices = self.downsample_modules[i](
                x, padding=padding_mask, mask_indices=mask_indices
            )

        residual = self.middle_encoder(x, padding_mask=padding_mask, layer=None)[0]
        x = x + residual
        res_outputs.append(x)

        # process decoders
        # The encoder has (self.label_nums - 1) blocks
        padding_masks.append(padding_mask)
        multi_mask_indices.append(mask_indices)
        residuals.reverse()  # NOTE(jiatong): reverse res_output to match corresponding input
        for i in range(self.label_nums - 1):
            x, padding_mask, mask_indices = self.upsample_modules[
                self.label_nums - 2 - i
            ](x, padding=padding_mask, mask_indices=mask_indices)
            x, _ = self.decoders[i](x, padding_mask=padding_mask, layer=None)
            x, padding_mask = align_size_sum(x, padding_mask, residuals[i])
            res_outputs.append(x)
            padding_masks.append(padding_mask)
            multi_mask_indices.append(mask_indices)

        # NOTE(jiatong): need reverse of target list to allow matched target-representation
        res_outputs.reverse()
        padding_masks.reverse()
        multi_mask_indices.reverse()
        if target_list is not None:
            new_target_list = []
            for i in range(self.label_nums):
                if self.use_single_target:
                    res_outputs[i], reformat_target_list = self.forward_targets(
                        res_outputs[i], target_list[0], self.feat2tar_ratios[i]
                    )
                    new_target_list.append(reformat_target_list)
                else:
                    if target_list[i] is not None:
                        res_outputs[i], reformat_target_list = self.forward_targets(
                            res_outputs[i], target_list[i], self.feat2tar_ratios[i]
                        )
                        new_target_list.append(reformat_target_list)
                    else:
                        # Append a None target list then it won't be used to calculate loss
                        new_target_list.append(None)
                if padding_masks[i] is not None:
                    padding_masks[i] = self.forward_padding_mask(
                        res_outputs[i], padding_masks[i]
                    )
                if multi_mask_indices[i] is not None:
                    multi_mask_indices[i] = self.forward_padding_mask(
                        res_outputs[i], multi_mask_indices[i]
                    )


        if features_only:
            # NOTE(jiatong): need to reverse back
            res_outputs.reverse()
            return {
                "x": res_outputs,
                "padding_mask": padding_masks[0],
                "features": features,
            }

        def compute_pred(proj_x, target, label_embs):
            # compute logits for the i-th label set
            y = torch.index_select(label_embs, 0, target.long())
            negs = label_embs.unsqueeze(1).expand(-1, proj_x.size(0), -1)
            if self.target_glu:
                y = self.target_glu(y)
                negs = self.target_glu(negs)
            # proj_x: (S, D)
            # y: (S, D)
            # negs: (Neg, S, D)
            return self.compute_nce(proj_x, y, negs)

        logit_m_list, logit_u_list = [], []
        for j in range(self.label_nums):
            if new_target_list[j] is None:
                continue  # skip empty targets
            label_embs_list = self.label_embs_concat[j].split(
                [self.multires_classes[j]], 0
            )
            # set the variables (after the set, the procedure is the same as hubert)
            # all the elements are list with only one element (to simulate the normal hubert process)
            x = res_outputs[j]
            target = new_target_list[j]
            padding_mask = padding_masks[j]
            mask_indices = multi_mask_indices[j]
            final_proj = self.final_projs[j]

            if not self.skip_masked:
                masked_indices = torch.logical_and(~padding_mask, mask_indices)
                proj_x_m = final_proj(x[masked_indices])
                logit_m_list.append(
                    compute_pred(proj_x_m, target[masked_indices], label_embs_list[0])
                )
            else:
                logit_m_list.append(None)

            if not self.skip_nomask:
                nomask_indices = torch.logical_and(~padding_mask, ~mask_indices)
                proj_x_u = final_proj(x[nomask_indices])
                logit_u_list.append(
                    compute_pred(proj_x_u, target[nomask_indices], label_embs_list[0])
                )
            else:
                logit_u_list.append(None)

            # if we only want one prediction, we can exit now
            if self.predictor_head_num == 1:
                break

        result = {
            "logit_m_list": logit_m_list,
            "logit_u_list": logit_u_list,
            "padding_mask": padding_mask,
            "features_pen": features_pen,
        }
        return result
output layer is 1-based
fairseq/models/multires_hubert/multires_hubert.py