src/speech_decoding/models/v14_converged_v3/objective.py

Selected training excerpt · lines 902–911 · source SHA-256 5718458b537d

902        f_max = max(nb for nb, _ in PER_BAND_SPECS)903        outs: list[Tensor] = []904        for head, (nb, _) in zip(self.mae_heads, PER_BAND_SPECS):905            o = head(h)  # (B, total, F_b)906            if nb < f_max:907                o = F.pad(o, (0, f_max - nb))908            outs.append(o)  # (B, total, F_MAX)909        stack = torch.stack(outs, dim=-1)  # (B, total, F_MAX, N_BANDS)910        onehot = F.one_hot(grid.band, len(PER_BAND_SPECS)).to(h.dtype)  # (total, N_BANDS)911        return (stack * onehot[None, :, None, :]).sum(-1)  # (B, total, F_MAX)