pytorch中报错IndexError: index 5 is out of bounds for dimension 0 with size 5

最近用pytorh复现yolov1时遇到上面的报错,想了很久最后查到是因为在Dataloader加载数据集时，如果不丢弃后面多余的batch，那么会出现这样的报错。
举个例子：

torch.Size([16, 3, 448, 448]) torch.Size([16, 7, 7, 30]
torch.Size([16, 3, 448, 448]) torch.Size([16, 7, 7, 30]
torch.Size([16, 3, 448, 448]) torch.Size([16, 7, 7, 30]
torch.Size([16, 3, 448, 448]) torch.Size([16, 7, 7, 30]
torch.Size([16, 3, 448, 448]) torch.Size([16, 7, 7, 30]
torch.Size([16, 3, 448, 448]) torch.Size([16, 7, 7, 30]
torch.Size([5, 3, 448, 448]) torch.Size([5, 7, 7, 30]
IndexError: index 5 is out of bounds for dimension 0 with size 5


# 上面是我的报错：输出的信息是print(images.size(), labels.size())

# 我发现最后的一个 batch_size 不是我设置的 batch_size 的大小，
# 然后经过思考我发现在我的训练集中有2501个数据，我的batch_size=16。那么2501/16余数刚好为5。

然后我又检查了Dataloader中的参数：

dataset：Dataset类型，从其中加载数据 
batch_size：int，可选。每个batch加载多少样本 
shuffle：bool，可选。为True时表示每个epoch都对数据进行洗牌 
sampler：Sampler，可选。从数据集中采样样本的方法。 
num_workers：int，可选。加载数据时使用多少子进程。默认值为0，表示在主进程中加载数据。 
collate_fn：callable，可选。 
pin_memory：bool，可选 
drop_last：bool，可选。True表示如果最后剩下不完全的batch,丢弃。False表示不丢弃。

最后我设置drop_last=True，代码正常运行。