KFrank:

F.one_hot(x, num_classes=4)


Thank you, this is exactly what I was looking for!
I have a follow up question though:
When I use this approach in my custom loss, I get the following error:
RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn

I guess the auto_grad function has troubles with the one_hot function.
any ideas on how to solve this?
Thanks in advance!