How to create a custom loss function in PyTorch

Loss functions define the scalar objective a PyTorch model tries to minimize. Built-in criteria cover common tasks, but domain-specific penalties often need extra weighting, masking, or piecewise math that still has to remain inside autograd.

Subclassing torch.nn.Module keeps the custom criterion usable like a built-in loss. Instantiate it once, call it with predictions and targets, and pass the returned scalar tensor to loss.backward().

A small CPU smoke script can prove the criterion before it moves into a training project. In the smoke workload, under-predictions receive a larger Huber-style penalty, then backward() fills a gradient and optimizer.step() changes a model weight.

1. Create a smoke script with a custom loss module.

   custom_loss_demo.py

   import torch
   from torch import nn
    
    
   class WeightedHuberLoss(nn.Module):
       def __init__(self, delta=0.5, under_prediction_weight=2.0):
           super().__init__()
           self.delta = delta
           self.under_prediction_weight = under_prediction_weight
    
       def forward(self, prediction, target):
           error = prediction - target
           abs_error = error.abs()
           quadratic = torch.minimum(abs_error, torch.full_like(abs_error, self.delta))
           linear = abs_error - quadratic
           huber = 0.5 * quadratic.pow(2) + self.delta * linear
           weights = torch.where(
               error < 0,
               torch.full_like(error, self.under_prediction_weight),
               torch.ones_like(error),
           )
           return (weights * huber).mean()
    
    
   torch.manual_seed(19)
    
   features = torch.tensor(
       [
           [0.0, 0.5],
           [1.0, -0.5],
           [2.0, 1.0],
           [-1.0, 1.5],
       ],
       dtype=torch.float32,
   )
   targets = torch.tensor([[0.4], [1.0], [2.4], [-0.2]], dtype=torch.float32)
    
   model = nn.Sequential(
       nn.Linear(2, 4),
       nn.Tanh(),
       nn.Linear(4, 1),
   )
   loss_fn = WeightedHuberLoss(delta=0.5, under_prediction_weight=2.5)
   optimizer = torch.optim.SGD(model.parameters(), lr=0.05)
    
   optimizer.zero_grad(set_to_none=True)
   prediction = model(features)
   loss = loss_fn(prediction, targets)
   loss.backward()
    
   gradient_norms = [
       parameter.grad.detach().norm()
       for parameter in model.parameters()
       if parameter.grad is not None
   ]
   gradient_norm = torch.linalg.vector_norm(torch.stack(gradient_norms), 2)
   first_layer_grad_present = model[0].weight.grad is not None
    
   first_layer_before = model[0].weight.detach().clone()
   optimizer.step()
   weight_delta = (model[0].weight.detach() - first_layer_before).abs().max()
    
   print(f"torch_version={torch.__version__}")
   print(f"custom_loss={loss.item():.6f}")
   print(f"gradient_norm={gradient_norm.item():.6f}")
   print(f"first_layer_grad_present={first_layer_grad_present}")
   print(f"optimizer_step_changed_weight={weight_delta.item() > 0}")

   forward() uses only torch tensor operations, so autograd can trace the loss back through the model. torch.where() applies the larger penalty only where the prediction is below the target.

2. Run the custom loss smoke script.

   $ python custom_loss_demo.py
   torch_version=2.12.1+cu130
   custom_loss=1.338794
   gradient_norm=1.481275
   first_layer_grad_present=True
   optimizer_step_changed_weight=True

3. Check the gradient and update lines.

   gradient_norm=1.481275 and first_layer_grad_present=True confirm loss.backward() reached model parameters. optimizer_step_changed_weight=True confirms the gradient produced an optimizer update.

4. Copy the loss class into the project code.

   class WeightedHuberLoss(nn.Module):
       def __init__(self, delta=0.5, under_prediction_weight=2.0):
           super().__init__()
           self.delta = delta
           self.under_prediction_weight = under_prediction_weight
    
       def forward(self, prediction, target):
           error = prediction - target
           abs_error = error.abs()
           quadratic = torch.minimum(abs_error, torch.full_like(abs_error, self.delta))
           linear = abs_error - quadratic
           huber = 0.5 * quadratic.pow(2) + self.delta * linear
           weights = torch.where(
               error < 0,
               torch.full_like(error, self.under_prediction_weight),
               torch.ones_like(error),
           )
           return (weights * huber).mean()

5. Instantiate the custom loss beside the optimizer.

   model = MyModel()
   optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
   loss_fn = WeightedHuberLoss(delta=0.5, under_prediction_weight=2.5)

6. Use the custom loss in the training step.

   for batch_features, batch_targets in train_loader:
       optimizer.zero_grad(set_to_none=True)
       predictions = model(batch_features)
       loss = loss_fn(predictions, batch_targets)
       loss.backward()
       optimizer.step()

   Keep predictions, batch_targets, and any tensors created inside forward() on the same device and dtype.
   Related: How to zero gradients in PyTorch

7. Keep non-tensor conversions out of forward().

   Do not call .detach(), .item(), or numpy() on values that still need gradients before returning the loss. Those conversions remove data from the autograd graph and can leave model parameters without gradients.

8. Remove the temporary smoke script.

   $ rm custom_loss_demo.py