drop_last when loading the dataset is necessary.

[yyotta]

The number of images in the mnist dataset is 60,000. At this time, if you set the batch size to 100 or any other value that can make 60,000 be divisible, that is OK. However, if you set other batch sizes that are commonly used, such as 256, 512, etc., it may lead to the wrong dimension of the last batch (96) and noise data (batch-size). So drop_last should be set when loading the dataset.
That is:
train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=bs, shuffle=True)
-> train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=bs, shuffle=True, drop_last=True)