Thank you, @Alexey_Demyanchuk
I got an error when trying your code, but then I changed it a bit and it worked like that:
def target_to_oh(target):
    NUM_CLASS = 3  # hard code here, can do partial
    one_hot = torch.eye(NUM_CLASS)[target]
    return one_hot