Introduction

Turning a model output tensor into a one-hot tensor is a common classification step. The usual approach is simple: find the winning class index, then encode that index as a one-hot vector. The important detail is that this only fits single-label classification, not every prediction problem.

Start With The Prediction Meaning

Before writing code, decide what the tensor represents. If each row means "exactly one class is correct," then argmax followed by one-hot encoding is appropriate. If several classes can be correct at the same time, forcing one winner is wrong.

For a standard multi-class output shaped like [batch_size, num_classes], the class dimension is usually axis=1.

1import tensorflow as tf
2
3scores = tf.constant([
4    [0.2, 1.7, 0.1],
5    [2.5, 0.3, 0.4],
6], dtype=tf.float32)
7
8predicted_index = tf.argmax(scores, axis=1)
9print(predicted_index.numpy())

This returns the index of the largest value in each row. That index is the class prediction.

Convert Indices To One-Hot

Once you have class indices, tf.one_hot creates indicator vectors with the correct depth.

1import tensorflow as tf
2
3scores = tf.constant([
4    [0.2, 1.7, 0.1],
5    [2.5, 0.3, 0.4],
6], dtype=tf.float32)
7
8predicted_index = tf.argmax(scores, axis=1)
9one_hot = tf.one_hot(predicted_index, depth=3)
10
11print(one_hot.numpy())

The output is:

[[0. 1. 0.]
 [1. 0. 0.]]

This is the normal post-processing pattern for classification predictions.

Logits And Probabilities Are Treated The Same

People often ask whether the model output must be softmax probabilities first. It does not. argmax only cares about which class is largest, so raw logits and normalized probabilities produce the same winner.

1import tensorflow as tf
2
3logits = tf.constant([[2.0, 0.5, 0.2]], dtype=tf.float32)
4probabilities = tf.nn.softmax(logits)
5
6print(tf.argmax(logits, axis=1).numpy())
7print(tf.argmax(probabilities, axis=1).numpy())

Both calls return the same class index. So if all you need is one-hot predicted classes, applying softmax first is optional.

Handle Batches And Sequence Outputs Correctly

The main source of bugs is using the wrong axis. For sequence models or token classifiers, the class dimension is often the last one rather than the second one.

1import tensorflow as tf
2
3sequence_scores = tf.random.uniform((2, 4, 3))
4predicted_index = tf.argmax(sequence_scores, axis=2)
5one_hot = tf.one_hot(predicted_index, depth=3)
6
7print(predicted_index.shape)
8print(one_hot.shape)

Here the input shape is [batch, time, classes]. The result after argmax is [batch, time], and after tf.one_hot it becomes [batch, time, classes].

A practical rule is: run argmax over the class axis, then use the number of classes as the one-hot depth.

Multi-Label Outputs Need Thresholding Instead

If the model can predict several classes as true at once, argmax discards valid information. In that case, apply a threshold per class.

1import tensorflow as tf
2
3multi_label_scores = tf.constant([
4    [0.9, 0.2, 0.8],
5    [0.1, 0.7, 0.6],
6], dtype=tf.float32)
7
8multi_hot = tf.cast(multi_label_scores >= 0.5, tf.float32)
9print(multi_hot.numpy())

This output is not one-hot in the strict sense because more than one position can be 1. That is correct for multi-label problems.

Use The Same Idea In Keras Prediction Code

In a Keras workflow, you often predict on a batch and convert the result immediately.

1import numpy as np
2import tensorflow as tf
3
4model = tf.keras.Sequential([
5    tf.keras.layers.Input(shape=(4,)),
6    tf.keras.layers.Dense(8, activation="relu"),
7    tf.keras.layers.Dense(3)
8])
9
10inputs = np.array([
11    [1.0, 0.2, 0.1, 0.0],
12    [0.0, 0.9, 0.3, 0.4],
13], dtype=np.float32)
14
15scores = model(inputs, training=False)
16predicted_index = tf.argmax(scores, axis=1)
17predicted_one_hot = tf.one_hot(predicted_index, depth=3)
18
19print(predicted_one_hot.numpy())

This is useful when downstream logic expects hard class assignments instead of floating-point scores.

PyTorch Has The Same Pattern

The same reasoning applies outside TensorFlow.

1import torch
2
3scores = torch.tensor([[0.2, 1.7, 0.1], [2.5, 0.3, 0.4]])
4predicted_index = torch.argmax(scores, dim=1)
5predicted_one_hot = torch.nn.functional.one_hot(predicted_index, num_classes=3)
6
7print(predicted_one_hot)

So the concept is framework-independent: pick class indices first, then encode them.

Common Pitfalls

• Using argmax on a multi-label output where several classes may be correct.
• Choosing the wrong axis for tensors that have time steps, channels, or other extra dimensions.
• Passing a depth value that does not match the real number of classes.
• Expecting softmax to be required before argmax for single-label prediction.
• Converting to one-hot too early when later code still needs confidence scores.

Summary

• For single-label classification, use argmax to get class indices and tf.one_hot to encode them.
• Raw logits and probabilities produce the same predicted class order.
• The class axis matters more than the tensor rank.
• Multi-label outputs should be thresholded, not collapsed with argmax.
• The same pattern works in TensorFlow, Keras, and PyTorch.