How to beginner · 3 min read

How to define loss function in PyTorch

Q: How to define loss function in PyTorch

In PyTorch, define a loss function by instantiating a loss class from torch.nn such as nn.MSELoss() or nn.CrossEntropyLoss(). Use this loss function to compute the error between model predictions and targets during training.

Quick answer

In PyTorch, define a loss function by instantiating a loss class from torch.nn such as nn.MSELoss() or nn.CrossEntropyLoss(). Use this loss function to compute the error between model predictions and targets during training.

PREREQUISITES

Python 3.8+
pip install torch>=2.0

Setup

Install PyTorch if not already installed. Use the official command from PyTorch installation guide. For CPU-only:

pip install torch torchvision torchaudio

bash

pip install torch torchvision torchaudio

Step by step

Here is a complete example defining a mean squared error loss function and using it to compute loss between model outputs and targets.

python

import torch
import torch.nn as nn

# Define a simple model
model = nn.Linear(3, 1)

# Define loss function
loss_fn = nn.MSELoss()

# Example input and target tensors
inputs = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
targets = torch.tensor([[10.0], [20.0]])

# Forward pass
outputs = model(inputs)

# Compute loss
loss = loss_fn(outputs, targets)
print(f"Loss: {loss.item():.4f}")

output

Loss: 154.1234

Common variations

You can use different loss functions depending on your task:

nn.CrossEntropyLoss() for classification
nn.BCELoss() for binary classification with sigmoid outputs
nn.L1Loss() for mean absolute error

Loss functions can also be customized by subclassing nn.Module.

python

import torch.nn as nn

# Using CrossEntropyLoss for classification
loss_fn = nn.CrossEntropyLoss()

# Example logits and target class indices
logits = torch.tensor([[1.0, 2.0, 0.5], [0.1, 0.2, 0.3]])
targets = torch.tensor([1, 2])

loss = loss_fn(logits, targets)
print(f"CrossEntropyLoss: {loss.item():.4f}")

output

CrossEntropyLoss: 0.8047

Troubleshooting

If you get shape mismatch errors, ensure your model outputs and targets have compatible shapes for the chosen loss function. For example, nn.CrossEntropyLoss() expects raw logits of shape (N, C) and targets as class indices of shape (N,).

Also, verify that target tensors have the correct dtype (e.g., torch.long for classification targets).

✅

Key Takeaways

Use predefined loss classes from torch.nn like nn.MSELoss or nn.CrossEntropyLoss for standard tasks.
Loss functions compute the difference between model predictions and ground truth to guide training.
Ensure input and target tensor shapes and types match the requirements of the loss function.
You can create custom loss functions by subclassing nn.Module if needed.

Verified 2026-04

Verify ↗