How do I one hot encode along a specific dimension using PyTorch?

user_9795226 0
Viewed 627

I have a tensor of size [3, 15, 136], where:

  • 3 is batch size
  • 15 - sequence length and
  • 136 is tokens

I want to one-hot my tensor using the probabilities in the tokens dimension (136). To do so I want to extract the tokens dimension for each letter in sequence length and put 1 to the largest possibility and mark all other tokens as 0.

one-hot-encoding tensor pytorch python
1 Answers

You can use PyTorch's one_hot function to achieve this:

import torch.nn.functional as F

t = torch.rand(3, 15, 136)

F.one_hot(t.argmax(dim=2), 136)
user_90676150
Related