With the latest release of the datasets library (2.19.0), the torch formatter now handles PIL objects of 2d arrays (e.g. grayscale images) differently.

Now when calling Dataset.with_format("torch"), every image tensor now has shape (1, H, W) for grayscale images. Previously, this was done by the unsqueeze(0) call in the subsequent step in the preprocessing pipeline, but calling unsqueeze is redundant for 2d image datasets (stored as PIL objects).

You can verify this by running:

from datasets import load_dataset

dataset = load_dataset("fashion_mnist", split="train")
transformed_dataset = dataset.with_format("torch")
transformed_dataset[:2]["image"].shape
#torch.Size([2, 1, 28, 28])

Previously, this would have returned

#torch.Size([2, 28, 28])

which required the call to unsqueeze.