def forward_scriptable(
        self,
        encoder_in,
        encoder_padding_mask,
        return_all_hiddens: bool = False,
        tgt_layer=None,
    ):
        """
        Args:
            src_tokens (LongTensor): tokens in the source language of shape
                `(batch, src_len)`
            src_lengths (torch.LongTensor): lengths of each source sentence of
                shape `(batch)`
            return_all_hiddens (bool, optional): also return all of the
                intermediate hidden states (default: False).
            token_embeddings (torch.Tensor, optional): precomputed embeddings
                default `None` will recompute embeddings

        Returns:
            dict:
                - **encoder_out** (Tensor): the last encoder layer's output of
                  shape `(src_len, batch, embed_dim)`
                - **encoder_padding_mask** (ByteTensor): the positions of
                  padding elements of shape `(batch, src_len)`
                - **encoder_embedding** (Tensor): the (scaled) embedding lookup
                  of shape `(batch, src_len, embed_dim)`
                - **encoder_states** (List[Tensor]): all intermediate
                  hidden states of shape `(src_len, batch, embed_dim)`.
                  Only populated if *return_all_hiddens* is True.
        """
        if self.no_freeze_encoder_layer is not None:
            ft = self.freeze_encoder_updates <= self.num_updates
        else:
            ft = True
        with torch.no_grad() if not ft else contextlib.ExitStack():
            # compute padding mask
            if not self.use_sent_enc_layer:
                has_pads = encoder_in.device.type == "xla" or encoder_padding_mask.any()

            if not self.layer_norm_first:
                encoder_in = self.layer_norm(encoder_in)

            encoder_in = self.dropout_module(encoder_in)

            # B x T x C -> T x B x C
            x = encoder_in.transpose(0, 1)

            encoder_states = []

            if return_all_hiddens:
                encoder_states.append(x)

            ## relative position embedding
            if self.args.relative_position_embedding:
                x_len = x.shape[0]
                pos_seq = torch.arange(0, x_len).long().to(x.device)
                pos_seq = pos_seq[:, None] - pos_seq[None, :]
                pos_k, pos_v = self.pos_emb(pos_seq)
            else:
                pos_k = None

        # encoder layers
        r = None
        d = None
        for i, layer in enumerate(self.layers):
            dropout_probability = np.random.random()

            with torch.no_grad() if (not ft) and i not in self.no_freeze_encoder_layer else contextlib.ExitStack():
                if not self.training or (dropout_probability > self.encoder_layerdrop) or i == self.unb_enc_layer:
                    if self.use_sent_enc_layer:
                        x, _ = layer(x, self_attn_padding_mask=encoder_padding_mask, self_attn_mask=None, need_weights=False, pos_bias=pos_k)
                        # x, _ = layer(x, self_attn_padding_mask=encoder_padding_mask, need_weights=False, pos_bias=pos_k)
                    else:
                        x = layer(x, encoder_padding_mask=encoder_padding_mask if has_pads else None, attn_mask=None)
                        # x = layer(x, encoder_padding_mask=encoder_padding_mask if has_pads else None)
                if i == self.unb_enc_layer:
                    d = x

                if i == tgt_layer:
                    r = x
                    break

                if return_all_hiddens:
                    assert encoder_states is not None
                    encoder_states.append(x)

        with torch.no_grad() if not ft else contextlib.ExitStack():
            # Finally T x B x C
            if self.layer_norm_first:
                x = self.layer_norm(x.transpose(0, 1)).transpose(0, 1)

            if r is not None:
                x = r

        # The Pytorch Mobile lite interpreter does not supports returning NamedTuple in
        # `forward` so we use a dictionary instead.
        # TorchScript does not support mixed values so the values are all lists.
        # The empty list is equivalent to None.
        return {
            "encoder_out": [x],  # T x B x C
            "encoder_padding_mask": [encoder_padding_mask],  # B x T
            "encoder_states": encoder_states,  # List[T x B x C]
            "src_tokens": [],
            "decoder_input": [d],
        }

        Args:
            src_tokens (LongTensor): tokens in the source language of shape
                `(batch, src_len)`
            src_lengths (torch.LongTensor): lengths of each source sentence of
                shape `(batch)`
            return_all_hiddens (bool, optional): also return all of the
                intermediate hidden states (default: False).
            token_embeddings (torch.Tensor, optional): precomputed embeddings
                default `None` will recompute embeddings

        Returns:
            dict:
                - **encoder_out** (Tensor): the last encoder layer's output of
                  shape `(src_len, batch, embed_dim)`
                - **encoder_padding_mask** (ByteTensor): the positions of
                  padding elements of shape `(batch, src_len)`
                - **encoder_embedding** (Tensor): the (scaled) embedding lookup
                  of shape `(batch, src_len, embed_dim)`
                - **encoder_states** (List[Tensor]): all intermediate
                  hidden states of shape `(src_len, batch, embed_dim)`.
                  Only populated if *return_all_hiddens* is True.
        
speecht5/speecht5/models/modules/encoder.py