Introduction

The PyTorch equivalent of TensorFlow softmax_cross_entropy_with_logits is usually torch.nn.CrossEntropyLoss or torch.nn.functional.cross_entropy. The important detail is that PyTorch expects raw logits, not probabilities, and in the common case it expects targets as class indices rather than one-hot vectors.

That is the source of most confusion when people translate code from TensorFlow. The names are similar, but the target format and the default workflow are not identical.

The Direct PyTorch Equivalent

For standard multiclass classification with integer labels, use CrossEntropyLoss. It combines a numerically stable log-softmax step with negative log likelihood loss in one operation.

1import torch
2import torch.nn as nn
3
4logits = torch.tensor([
5    [2.1, 0.3, -1.2],
6    [0.1, 1.7, 0.2],
7], dtype=torch.float32)
8
9targets = torch.tensor([0, 1], dtype=torch.long)
10
11criterion = nn.CrossEntropyLoss()
12loss = criterion(logits, targets)
13print(loss.item())

Do not apply softmax before this loss. CrossEntropyLoss already handles the normalization internally and is more numerically stable than doing it in two separate steps.

What This Corresponds To in TensorFlow

In TensorFlow, softmax_cross_entropy_with_logits takes logits and a label distribution. In many examples, that distribution is one-hot encoded. In PyTorch, the closest everyday translation is to keep logits exactly as they are and convert one-hot labels into class indices.

1import torch
2import torch.nn.functional as F
3
4logits = torch.tensor([
5    [2.1, 0.3, -1.2],
6    [0.1, 1.7, 0.2],
7])
8
9one_hot = torch.tensor([
10    [1.0, 0.0, 0.0],
11    [0.0, 1.0, 0.0],
12])
13
14targets = one_hot.argmax(dim=1)
15loss = F.cross_entropy(logits, targets)
16print(loss.item())

If your labels are truly one-hot and represent a single correct class, converting with argmax is normally the right move.

When Targets Are Soft Labels

Sometimes the TensorFlow code is using label smoothing or a non one-hot target distribution. In that case, converting with argmax loses information. You need the loss against the full target distribution.

A manual PyTorch implementation is straightforward:

1import torch
2import torch.nn.functional as F
3
4logits = torch.tensor([
5    [2.1, 0.3, -1.2],
6    [0.1, 1.7, 0.2],
7], dtype=torch.float32)
8
9soft_targets = torch.tensor([
10    [0.8, 0.2, 0.0],
11    [0.1, 0.7, 0.2],
12], dtype=torch.float32)
13
14log_probs = F.log_softmax(logits, dim=1)
15loss = -(soft_targets * log_probs).sum(dim=1).mean()
16print(loss.item())

That formula is the closer semantic match when the TensorFlow code genuinely uses probability distributions as labels.

Why You Should Not Call softmax First

This is the most common translation bug. Developers see softmax_cross_entropy_with_logits, notice the word softmax, and explicitly apply softmax before the PyTorch loss. That is wrong for CrossEntropyLoss.

Wrong pattern:

probs = torch.softmax(logits, dim=1)
loss = F.cross_entropy(probs, targets)

Correct pattern:

loss = F.cross_entropy(logits, targets)

Passing probabilities instead of logits changes the math and can make optimization worse.

Binary Classification Is Slightly Different

If the TensorFlow code is for binary classification and the model outputs one logit per example, the better PyTorch equivalent is often BCEWithLogitsLoss, not CrossEntropyLoss.

1import torch
2import torch.nn as nn
3
4logits = torch.tensor([0.9, -1.4, 0.2])
5targets = torch.tensor([1.0, 0.0, 1.0])
6
7criterion = nn.BCEWithLogitsLoss()
8loss = criterion(logits, targets)
9print(loss.item())

The choice depends on the model output shape and whether the task is multiclass, binary, or multilabel.

Check Shapes and Dtypes

For ordinary CrossEntropyLoss usage:

• logits shape is usually (batch_size, num_classes)
• target shape is usually (batch_size,)
• target dtype should be integer class indices such as torch.long

Shape and dtype mistakes often produce cryptic runtime errors, so this is worth checking early.

Common Pitfalls

The biggest mistake is applying softmax before CrossEntropyLoss. Another is feeding one-hot labels directly into the standard multiclass loss when the code actually expects class indices. Developers also mis-handle soft labels by collapsing them with argmax even when the distribution itself matters. Finally, some binary problems are implemented with the wrong loss entirely and should have used BCEWithLogitsLoss from the start.

Summary

• The usual PyTorch equivalent is nn.CrossEntropyLoss on raw logits.
• Do not apply softmax before the loss.
• Convert one-hot labels to class indices when the problem is ordinary multiclass classification.
• If the TensorFlow code uses soft target distributions, compute the loss against log_softmax directly.
• For binary or multilabel setups, BCEWithLogitsLoss may be the real equivalent.