The test data of MNIST will contain 10000 samples.
If you are using a batch size of 64, you would get 156 full batches (9984 samples) and a last batch of 16 samples (9984+16=10000), so I guess you are only checking the shape of the last batch.
If you don’t want to use this last (smaller) batch, you can use drop_last=True in the DataLoader.