When using mnist or fashion-mnist these two datasets, the "transform" of the dataloader.py should be "transforms.Compose([transforms.Resize((input_size, input_size)), transforms.ToTensor(), transforms.Normalize(mean=([0.5]), std=([0.5]))])" ,instead of "transforms.Compose([transforms.Resize((input_size, input_size)), transforms.ToTensor(), transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))])" .
When using mnist or fashion-mnist these two datasets, the "transform" of the dataloader.py should be "transforms.Compose([transforms.Resize((input_size, input_size)), transforms.ToTensor(), transforms.Normalize(mean=([0.5]), std=([0.5]))])" ,instead of "transforms.Compose([transforms.Resize((input_size, input_size)), transforms.ToTensor(), transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))])" .