    def forward(self, x, mask, ref_emb, ref_mask):
        """
        input:
        x: (B, N, d)
        mask: (B, N), mask is 0
        ref_emb: (B, d, T')
        ref_mask: (B, T'), mask is 0

        output:
        dur_pred: (B, N)
        dur_pred_log: (B, N)
        dur_pred_round: (B, N)
        """

        input_ref_mask = ~(ref_mask.bool())  # (B, T')
        # print(input_ref_mask)

        x = x.transpose(1, -1)  # (B, N, d) -> (B, d, N)

        for idx, (conv, act, ln, dropout) in enumerate(self.conv):
            res = x
            # print(torch.min(x), torch.max(x))
            if idx % self.cross_attn_per_layer == 0:
                attn_idx = idx // self.cross_attn_per_layer
                attn, attn_ln, attn_drop = self.cattn[attn_idx]

                attn_res = y_ = x.transpose(1, 2)  # (B, d, N) -> (B, N, d)

                y_ = attn_ln(y_)
                # print(torch.min(y_), torch.min(y_))
                # print(torch.min(ref_emb), torch.max(ref_emb))
                y_, _ = attn(
                    y_,
                    ref_emb.transpose(1, 2),
                    ref_emb.transpose(1, 2),
                    key_padding_mask=input_ref_mask,
                )
                # y_, _ = attn(y_, ref_emb.transpose(1, 2), ref_emb.transpose(1, 2))
                # print(torch.min(y_), torch.min(y_))
                y_ = attn_drop(y_)
                y_ = (y_ + attn_res) / math.sqrt(2.0)

                x = y_.transpose(1, 2)

            x = conv(x)
            # print(torch.min(x), torch.max(x))
            x = act(x)
            x = ln(x.transpose(1, 2))
            # print(torch.min(x), torch.max(x))
            x = x.transpose(1, 2)

            x = dropout(x)

            if idx != 0:
                x += res

            if mask is not None:
                x = x * mask.to(x.dtype)[:, None, :]

        x = self.linear(x.transpose(1, 2))
        x = torch.squeeze(x, -1)

        dur_pred = x.exp() - 1
        dur_pred_round = torch.clamp(torch.round(x.exp() - 1), min=0).long()

        return {
            "dur_pred_log": x,
            "dur_pred": dur_pred,
            "dur_pred_round": dur_pred_round,
        }
input:
x: (B, N, d)
mask: (B, N), mask is 0
ref_emb: (B, d, T')
ref_mask: (B, T'), mask is 0

output:
dur_pred: (B, N)
dur_pred_log: (B, N)
dur_pred_round: (B, N)
modules/naturalpseech2/transformers.py