Thank you for your comments. I can’t make it work (I no tested more nn.BCEWithLogitsLoss, I focused in N=5 multi class segmentation). Images are [1,220,220] and  their mask [5,220,220]. Each channel is a binary image with values 0 and 1, 1s for the object of interest on the respective channel, and 0s for the background. I don’t know if this mask form is correct.
The error occurs in the line

loss = criterion(output, target) ,

during the first pass. The error seems related to data type, but If I am not mistaken, internal tensors work with float numbers (all images and masks are floats).
Here is my code:
model = UnetModel(n_filters=32, n_class=5, input_channels=1)
device = ‘cuda’
model = model.to(device)
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
model.train()
epochs = 20
itr = 1
p_itr = 50
total_loss = 0
train_loss = []
train_acc = []
for epoch in range(epochs):
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**
    loss.backward()  
    optimizer.step() 
            
    total_loss += loss.item()        
    if itr%p_itr == 0:
        pred = torch.argmax(output, dim=1)
        correct = pred.eq(target)
        acc = torch.mean(correct.float())
        print('[Epoch {}/{}] Iteration {} -> Train Loss: {:.4f}, Accuracy: {:.3f}'.format(epoch+1, epochs, itr, total_loss/p_itr, acc))
        train_loss.append(total_loss/p_itr)
        train_acc.append(acc)
        total_loss = 0            
    itr += 1

The error is
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