I am really thankful for fast replying I really need to finish this code so I am also sorry for asking so question~~
After I finished my dataset part I was started to train part like that basically:
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
model = UNet()
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)


for epoch in range(epochs):
    for i, (images, masks) in enumerate(train_loader, 1):
        print(images.size())
        for idx in range(images.size(3)):
            images = images[:, :, :, idx].unsqueeze(1).to(device)
            masks = masks[:, :, :, idx].unsqueeze(1).to(device)
            images = images.type(torch.FloatTensor)
            masks = images.type(torch.FloatTensor)
            print(images.shape)
            print(masks.shape)
            outputs = model(images)
            loss = criterion(outputs, masks)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            if (i) % 1 == 0:
                print(f"Epoch [{epoch + 1}/{epochs}], Step [{i}/{total_steps}], Loss: {loss.item():4f}")

But I couldn`t understand why it was connected to CrossEntropyLoss because it is giving this error from loss part : RuntimeError: only batches of spatial targets supported (3D tensors) but got targets of dimension: 4