How to make a 3D tensor into one_hot encoding
for example ,labels is a groundtruth, labels size is [4,224,224] where 4 means batchsize ,224 means h and w,  labels dtype is torch.int64, in my training code ,label's pixel from 0 to 255 .I use my network to Semantic segmentation

From your question, it seems like you would like 256 integer labels, from 0 to 255. This can be done by the following:

# let a be the (4, 224, 224) dim tensor
labels = np.stack([(a == i).int() for i in range(256)])
print(labels.shape) #(256, 4, 224, 224) dimensional one-hot encoding


Hope this helps!
