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.