You could do the following:
criterion = torch.nn.CrossEntropy(reduction='none')

This ensures that the function will return a loss value for each element. You could then multiply the weights to each loss element.
gt # Ground truth, format torch.long
pd # Network output
W # per-element weighting based on the distance map from UNet
loss = criterion(pd, gt)
loss = W*loss # Ensure that weights are scaled appropriately
loss = torch.sum(loss.flatten(start_dim=1), axis=0) # Sums the loss per image
loss = torch.mean(loss) # Average across a batch

I’ve observed that a weight of 0 could provide complications. It might be a good idea to add a small eps factor to W.