Based on the stacktrace it seems you are using a Dice loss, which fails in the one_hot call on the targets (or rather the scatter_ operation). Could you post the shapes of logit_map and y in this line of code:
loss = loss_function(logit_map, y)

(I see that you’ve already added the print statements so you are on the right path to debug it  )