Consider the following 6x6 tensor.
tst_tensor = tensor([[1, 1, 1, 2, 3, 1],
        [1, 2, 0, 0, 0, 0],
        [1, 0, 1, 0, 0, 0],
        [5, 5, 5, 5, 5, 5],
        [5, 5, 5, 0, 0, 0],
        [0, 2, 3, 0, 0, 0]])

I’d like to window it in 2x2 windows (or 4 x 1) as follows
tensor([[1, 1, 1, 2],
        [1, 2, 0, 0],
        [3,1, 0, 0] ...
        [0,0,0,0]])

However if I run the following code
batched_input= tst_tensor.unsqueeze(0).unsqueeze(0)
 unfold = nn.Unfold(kernel_size=(3,3),stride=3,dilation=1)
 windows = unfold(batched_input)

I get something of the right shape (1,9,4), but the values are not what I expect:
tensor([[[1., 2., 5., 5.],
         [1., 3., 5., 5.],
         [1., 1., 5., 5.],
         [1., 0., 5., 0.],
         [2., 0., 5., 0.],
         [0., 0., 5., 0.],
         [1., 0., 0., 0.],
         [0., 0., 2., 0.],
         [1., 0., 3., 0.]]])

Where am I going wrong?