[intro-to-pytorch-part7]Img load issue for help

[LeonMac]

Hello, for the "load image data" section, I nearly copy the code, but I always just get one image shown for both trainloader and testloader.

Below is my code (sorry I always have trouble pasting code in github, the code between `` can not be shown nicely--anyone can help to check where the problem could be?
I am sure I have been downloaded all the necessary image..

Thanks!

`root_dir = './Cat_Dog_data'
batch_sz = 32
max_batch_num = 3
train_transforms = transforms.Compose([transforms.RandomRotation(30),
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor()])

test_transforms = transforms.Compose([transforms.Resize(255),
transforms.CenterCrop(224),
transforms.ToTensor()])

train_data = datasets.ImageFolder(root_dir + '/train/', transform=train_transforms)
test_data = datasets.ImageFolder(root_dir + '/test/', transform=test_transforms)

trainloader = torch.utils.data.DataLoader(train_data, batch_size=batch_sz)
testloader = torch.utils.data.DataLoader(test_data, batch_size=batch_sz)

for i in range (2):
if i == 1:
images, labels = next(iter (trainloader))
else:
images, labels = next(iter (testloader) )
n = 4
fig, axes = plt.subplots(figsize=(10,n), ncols=n)

for ii in range(n):
    ax = axes[ii]
    helper.imshow(images[ii], ax=ax, normalize=False) #, auto_close=True)`