naplab / conv-tasnet Goto Github PK

    # input is the waveforms: (B, T) or (B, 1, T)
    # reshape and padding
    if input.dim() not in [2, 3]:
        raise RuntimeError("Input can only be 2 or 3 dimensional.")
    
    if input.dim() == 2:
        input = input.unsqueeze(1)
    # print(input.shape)
    batch_size = input.size(0)
    nsample = input.size(2)
    rest = self.win - (self.stride + nsample % self.win) % self.win
    # print(self.win,rest)
    if rest > 0:
        pad = Variable(torch.zeros(batch_size, 1, rest)).type(input.type())
        input = torch.cat([input, pad], 2)
    
    pad_aux = Variable(torch.zeros(batch_size, 1, self.stride)).type(input.type())
    input = torch.cat([pad_aux, input, pad_aux], 2)

    return input, rest`