I have an image segmentation (Learned Random Walker), but I would like to generalize it.
# Init the random walker modules
rw = RandomWalker(1000, max_backprop=True)

# Load data and init
target = init_image(imagem)
seeds = init_seeds(torch.zeros(1, 60, 59).long(), target)  
diffusivities = torch.zeros(1, 2, 60, 59, requires_grad=True)

# Init optimizer
optimizer = torch.optim.Adam([diffusivities], lr=0.5)

loss = torch.nn.NLLLoss()

# Main overfit loop
for it in range(iterations + 1):
    optimizer.zero_grad()

    # Diffusivities must be positive
    net_output = torch.sigmoid(diffusivities)

    # Random walker
    output = rw(net_output, seeds)

    # Loss and diffusivities update
    m = torch.nn.LogSoftmax(dim=1)
    output_log = [m(o) for o in output]

    # ls = loss(output_log, target)
    ls = loss(output_log[0], target[0])
    ls.backward()
    optimizer.step()

Given target of dimension [1, 1, 60, 59], output will be of dimension [1, nClass, 60, 59], that is, it will be the nClass probability matrices for each region.
But for example, with this same dimension, I have the image target in the range [0, 255] with 3 regions, NLLLoss does not allow me to do this, because I would have output with dimension [1, 3, 60, 59] and several pixels greater than 3.
Is there anything you can do to get around this, or some loss function that solves this problem?
Any help is welcome, thanks.