This part of the error suggests the problem is in your use of torchvision.transforms.Normalize.

–> 212         return F.normalize(tensor, self.mean, self.std, self.inplace)

The docs for that name are here: https://pytorch.org/docs/stable/torchvision/transforms.html#torchvision.transforms.Normalize.
That page describes its arguments as:


mean (sequence) – Sequence of means for each channel.

std (sequence) – Sequence of standard deviations for each channel.

In your code, you gave 3 means and 3 standard deviations to Normalize, which it will try to use for 3 different channels:
torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
Since your inputs only have 1 channel, you should only be passing lists of 1 mean and 1 standard deviation like this:
torchvision.transforms.Normalize(mean=[0.485], std=[0.229])