Torch.nn.functional.one_hot does not work as expected with ROCM

I installed the pytorch 2.1 with rocm 5.4.2.
When I tested the one_hot function, it returned this:
print(torch.nn.functional.one_hot(torch.tensor(0,device=‘cuda:0’), 2)) = [0,0]

Do you know why it behaved like this? Thank you in advance

This looks like a bug and given that you already have minimal reproducible example I would recommend opening an issue @ GitHub - ROCmSoftwarePlatform/pytorch: Tensors and Dynamic neural networks in Python with strong GPU acceleration