RuntimeError：对于4维权重[64、1、3、3]，需要4维输入，但是却得到了大小为[1、1、128、128、128]的5维输入

Question

使用此UNet实现：https://github.com/kilgore92/PyTorch-UNet

使用UNet对象的此实例：

model = UNet(n_channels=1,
             mode='3D',
             num_classes=1,
             use_pooling=True,
             )

每当我尝试使用dataLoader运行训练脚本时，都会出现此错误：

RuntimeError: Expected 4-dimensional input for 4-dimensional weight [64, 1, 3, 3], but got 5-dimensional input of size [1, 1, 128, 128, 128] instead

更多信息：

这是我的训练脚本，在输出= model.forward（images）处出错。

def train():
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
    num_steps_train = len(train_loader)
    
    print(num_steps_train)
    
    for epoch in range(epochs):
        print(' - training - ')
        for i, (images, masks) in enumerate(train_loader):
            images = images.to(device)
            masks = masks.to(device)
            outputs = model.forward(images)

当我打印出来时，图像和遮罩的大小为（1、1、128、128）。这是因为在数据集对象的getitem_中，我使用image = torch.reshape（image，shape =（1，128，128，128））将图像重塑为指定的大小。我尝试将其更改为image = torch.reshape（image，shape =（128，128，128）），但这也不起作用。

这是我的train_loader：

train_loader = DataLoader(dataset=Dataset(partition['orig'], partition['segment']), 
                          batch_size = batch_size, shuffle = True)

分区只是其中包含原始蒙版图像和细分蒙版图像的字典。

我似乎无法弄清楚。我的形状正在对齐，并且在需要时创建了新轴。