def segment_specs(y, seg_length=15, seg_hop=4, max_length=None):
    '''
    Segment a spectrogram into "seg_length" wide spectrogram segments.
    Instead of using only the frequency bin of the current time step, 
    the neighboring bins are included as input to the CNN. For example 
    for a seg_length of 7, the previous 3 and the following 3 frequency 
    bins are included.

    A spectrogram with input size [H x W] will be segmented to:
    [W-(seg_length-1) x C x H x seg_length], where W is the width of the 
    original mel-spec (corresponding to the length of the speech signal),
    H is the height of the mel-spec (corresponding to the number of mel bands),
    C is the number of CNN input Channels (always one in our case).
    '''
    if seg_length % 2 == 0:
        raise ValueError('seg_length must be odd! (seg_lenth={})'.format(seg_length))
    if not torch.is_tensor(y):
        y = torch.tensor(y)

    B, _, _ = y.size()
    for b in range(B):
        x = y[b,:,:]
        n_wins = x.shape[1]-(seg_length-1)
        # broadcast magic to segment melspec
        idx1 = torch.arange(seg_length)
        idx2 = torch.arange(n_wins)
        idx3 = idx1.unsqueeze(0) + idx2.unsqueeze(1)
        x = x.transpose(1,0)[idx3,:].unsqueeze(1).transpose(3,2)

        if seg_hop>1:
            x = x[::seg_hop,:]
            n_wins = int(np.ceil(n_wins/seg_hop))

        if max_length is not None:
            if max_length < n_wins:
                raise ValueError('n_wins {} > max_length {}. Increase max window length ms_max_segments!'.format(n_wins, max_length))
            x_padded = torch.zeros((max_length, x.shape[1], x.shape[2], x.shape[3]))
            x_padded[:n_wins,:] = x
            x = x_padded
        if b == 0:
            z = x.unsqueeze(0)
        else:
            z = torch.cat((z, x.unsqueeze(0)), axis = 0)
    B, n, c, f, t = z.size()
    z = z.view(B*n, c, f, t)
    return z

    Segment a spectrogram into "seg_length" wide spectrogram segments.
    Instead of using only the frequency bin of the current time step, 
    the neighboring bins are included as input to the CNN. For example 
    for a seg_length of 7, the previous 3 and the following 3 frequency 
    bins are included.

    A spectrogram with input size [H x W] will be segmented to:
    [W-(seg_length-1) x C x H x seg_length], where W is the width of the 
    original mel-spec (corresponding to the length of the speech signal),
    H is the height of the mel-spec (corresponding to the number of mel bands),
    C is the number of CNN input Channels (always one in our case).
    
train/speech_enhancement/models/mossformer_gan/discriminator.py