Thanks!
I’ll leave the function I’m using here in case anyone else can use it
def windowed_tensor(tensor, size, stride):
    """assumes a square 3D tensor"""
    tensor = tensor.unfold(1, size, stride)
    tensor = tensor.unfold(2, size, stride)
    tensor = torch.transpose(tensor, 0, 1)
    tensor = torch.transpose(tensor, 1, 2)
    tensor = tensor.contiguous()
    tensor = tensor.view(-1, 3, size, size)
    return tensor