pytorch stack widow in timeseries

pytorch_window_stack.py

def window_stack(x, window=4, pad=True):
    """
    Stack along a moving window of a pytorch timeseries
    
    Inputs: 
        tensor of dims (batches/time, channels)
        pad: if true the left side will be padded to let the output match
    Outputs: 
        if pad=True: a tensor of size (batches, channels, window)
        else: tensor of size (batches-window, channels, window)
    """
    x = x.transpose(0,1)
    if pad:
        x = F.pad(x, (window, 0)+(0,0)*(len(x.size())-1))
    print(x.size(1))
    x = torch.stack([x[:, i:i+window] for i in range(x.size(1)-window)])
    return x

input = torch.range(0,19).view((2,10)).transpose(1,0) # batch/time, channels
print(input)
x = window_stack(input, window=4, pad=False)
x