Here’s ReshapeTransform I wrote:
class ReshapeTransform:
    def __init__(self, new_size):
        self.new_size = new_size

    def __call__(self, img):
        return torch.reshape(img, self.new_size)

You can use this for MNIST like this:
        mnist_transforms=[
            transforms.ToTensor(),
            transforms.Normalize((0.1307,), (0.3081,)),
            mnist_transforms.append(ReshapeTransform((-1,)))]
        train_ds = datasets.MNIST(train=True, download=True, 
            transform=transforms.Compose(mnist_transforms));