torch中Dataloader类加载数据，同样的数据在训练集中返回的纬度与在测试集中返回的纬度不一样

问题遇到的现象和发生背景

在使用Torch进行训练时发现Dataloader在训练和测试时返回的数据不是同一个纬度的。我喂入的数据形状是(5,3,224,224)，由于batch_size我设置了2，返回的应该是（2,5,3,224,224)，训练集返回正常，但是测试集返回的形状仍然是(5,3,224,224)

问题相关代码，请勿粘贴截图

train_data = MyDataset('train.txt', transform=transform, loader=MyLoader)
train_loader = DataLoader(dataset=train_data, batch_size=2, shuffle=True)

test_data = MyDataset('train.txt', transform = transform, loader=MyLoader)
test_loader = DataLoader(test_data,batch_size=2, shuffle=False)

运行结果及报错内容

运行结果为test_imgs(5,3,224,224)，train_imgs(2,5,3,224,224)。由于我网络的定义，期望的输入形状是五维的所以会报错

我的解答思路和尝试过的方法

我用unsqueeze方法增加了一个纬度就可以了，但是不清楚为什么一个返回五维，一个返回四维

我想要达到的结果

我是SB，我写错了，for循环里应该引用test_loader