batch size one yields squeezed dimensions when outputting one dimensional arrays

[lengstrom]

cifar betons made using the tutorial.

from pathlib import Path
import numpy as np
import torch as ch
import torchvision

from ffcv.fields import IntField, RGBImageField
from ffcv.fields.decoders import IntDecoder, SimpleRGBImageDecoder
from ffcv.loader import Loader, OrderOption
from ffcv.pipeline.operation import Operation
from ffcv.transforms import RandomHorizontalFlip, Cutout, \
    RandomTranslate, Convert, ToDevice, ToTensor, ToTorchImage
from ffcv.transforms.common import Squeeze
from ffcv.writer import DatasetWriter

def make_loader(part, indices, batch_size):
    path = {
        'train': '/mnt/xfs/home/engstrom/store/cifar_datasets/cifar_train.beton',
        'test': '/mnt/xfs/home/engstrom/store/cifar_datasets/cifar_test.beton'
    }[part]

    CIFAR_MEAN = [125.307, 122.961, 113.8575]
    CIFAR_STD = [51.5865, 50.847, 51.255]
    loaders = {}

    label_pipeline: List[Operation] = [IntDecoder(), ToTensor(), ToDevice(ch.device('cuda:0')), Squeeze()]
    image_pipeline: List[Operation] = [SimpleRGBImageDecoder()]
    image_pipeline.extend([
        ToTensor(),
        ToDevice(ch.device('cuda:0'), non_blocking=True),
        ToTorchImage(),
        Convert(ch.float32),
        torchvision.transforms.Normalize(CIFAR_MEAN, CIFAR_STD),
    ])

    ordering = OrderOption.SEQUENTIAL
    loader = Loader(path, batch_size=batch_size, num_workers=2,
                    order=ordering, drop_last=False,
                    pipelines={'image': image_pipeline, 'label': label_pipeline},
                    indices=indices)
    return loader

if __name__ == '__main__':
    loader1 = make_loader('test', range(20), 1)
    loader2 = make_loader('test', range(20), 2)
    print('loader bs 1')
    for tup in loader1:
        for t in tup:
            print(t.shape)
        break

    print('loader bs 2')
    for tup in loader2:
        for t in tup:
            print(t.shape)
        break

yields

loader bs 1
torch.Size([1, 3, 32, 32])
torch.Size([])
loader bs 2
torch.Size([2, 3, 32, 32])
torch.Size([2])

but it should yield:

loader bs 1
torch.Size([1, 3, 32, 32])
torch.Size([1]) <-- here
loader bs 2
torch.Size([2, 3, 32, 32])
torch.Size([2])

[wouterzwerink]

Is that not simply due to the Squeeze transform? Perhaps that transform should be able to take in dimension arguments.