ミニバッチを取り出して学習できない

前提・実現したいこと

色々なソースコードをみて、機械学習をしようとしています。
そのソースコードでは、dataloaderをfor分でイテレーションを行っています。
しかし、このような僕のソースコードでは、なぜか、ミニバッチサイズの回数分を学習をしないです。

データーローダーからミニバッチを取り出すループを実現したいです。

少し説明しづらくて、わかりにくいと思うのですが、教えて頂けるとうれしいです。

version
python　3.6
pytorch 1.7.1
torchvision 0.8.3

発生している問題

import torch
import torchvision
from torch.utils import data
from torchvision import transforms
#%%

train_dataset = torchvision.datasets.ImageFolder(
    root='/Users/name/DATASETS',
    transform= transforms.Compose([
        transforms.Resize(size=64),
        transforms.CenterCrop(size=64),
        transforms.ToTensor(),
        transforms.Normalize(mean=(0.5,0.5,0.5), std=(0.5,0.5,0.5))
    ])
)
train_loader = torch.utils.data.DataLoader(
    dataset=train_dataset,
    batch_size=3,
    shuffle=True
)

#%%

for epoch in range(1000):
    for i,(img, label) in enumerate(train_loader):
        print(i)
    print('Epoch:', epoch)

>>> (結果)
1
2
3
4
5
6
.
.
.
（永遠続く。）
ここで学習したいのですが、バッチサイズ回行ってくれません。

こうなってほしいソースコード

#%%
for epoch in range(1000):
    for i,(img, label) in enumerate(train_loader):
        print(i)
    print('Epoch:', epoch)
>>> (結果)
1
2
3
Epoch:0
1
2
3
Epoch:1
1
.
.
.
.