I am a new user of Pytorch. I am adapting the Unet segmentation model, but I have an error in the evaluation of the Cross Entropy Loss  function during training.
I used  torch.utils.data.Dataset to build a specific dataset
train_data = DataLoaderSegmentation(train_path, mode=‘train’)
train_loader = DataLoader(train_data, batch_size = 4, shuffle=True, num_workers=4)
which works fine and load both the image and its mask in the format:
[batch,channel,W,H] , like  [4,1,220,220], for both images and masks. I put channel to 1 since for background and foreground segmentation.
Parameters for train:
model = Unet(n_filters=32, n_class=1, input_channels=1)
device = ‘cuda’
LR = 0.001
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=LR)
model = model.to(device)
During training I have an error on the marked line:
for data, target in train_loader:
data = data.to(device)
target = target.to(device)
optimizer.zero_grad()
output = model(data)
loss = criterion(output, target) ############## ERROR
for data, target in train_loader:
data = data.to(device)
target = target.to(device)
optimizer.zero_grad()
output = model(data)
loss = criterion(output, target) ######## ERROR

File “C:\Anaconda3\envs\envTorch\lib\site-packages\torch\nn\functional.py”, line 1840, in nll_loss
**    ret = torch._C._nn.nll_loss2d(input, target, weight, _Reduction.get_enum(reduction), ignore_index)**
RuntimeError: Expected object of scalar type Long but got scalar type Float for argument #2 ‘target’ in call to _thnn_nll_loss2d_forward

I see that  target and output have the same dim [4,1,220,220] but  the criterion(output, target) trow error.
Sorry, but after read some answer, I cannot resolve it yet.