When implementing a custom loss function as a class in PyTorch, which base class should it inherit from to integrate properly with the training loop?
Select an answer to reveal the explanation.
Short Explanation
Here's the pattern PyTorch wants you to follow: a custom loss, just like a custom layer or a custom model, should inherit from torch.nn.Module. You define your loss computation inside forward(), and because everything inside is built from standard differentiable tensor operations, autograd can trace through it automatically to compute gradients — no extra plumbing needed. That's why torch.nn.Module is correct. Dataset is for wrapping and indexing your data samples, totally unrelated to computing a loss value. Optimizer is about updating weights after gradients exist, not computing the loss that produces those gradients in the first place. And torch.autograd.Function is a lower-level tool for defining custom forward/backward operations manually, which is overkill and unnecessary for a standard loss built from existing differentiable ops.
Full Explanation
In PyTorch, custom loss functions implemented as classes should inherit from torch.nn.Module, mirroring the same pattern used for defining custom model architectures and layers. By subclassing nn.Module and implementing the loss computation inside the forward() method using standard differentiable PyTorch tensor operations, the loss automatically integrates with PyTorch's autograd system: calling loss.backward() will correctly trace back through every operation performed inside forward() to compute gradients with respect to any tensors that require them (such as model parameters), with no additional manual work required. This is exactly how PyTorch's own built-in losses, such as nn.CrossEntropyLoss and nn.MSELoss, are implemented internally.
torch.utils.data.Dataset is the base class for defining how individual data samples are accessed (via __getitem__ and __len__), used in conjunction with DataLoader to feed batches into training. It has no relationship to defining a loss computation and would not provide any of the necessary machinery for computing or backpropagating a loss value.
torch.optim.Optimizer is the base class for optimization algorithms such as SGD and Adam, which take computed gradients and apply parameter updates. An optimizer consumes the output of the backward pass (the gradients); it does not compute the loss itself, so inheriting from it for a custom loss class would be a fundamental misunderstanding of the training loop's separation of concerns.
torch.autograd.Function is a lower-level API used when a developer needs to manually define custom forward and backward computation logic, typically for operations that are not expressible through standard differentiable PyTorch operations (for example, wrapping a non-differentiable external library call and supplying a hand-written gradient). For the vast majority of custom losses, which can be expressed using existing differentiable tensor operations, this level of manual control is unnecessary overhead; nn.Module with a standard forward() implementation is sufficient and is the conventional, recommended approach.
Memory aid: in PyTorch, both models and losses are commonly built as nn.Module subclasses because both are fundamentally 'differentiable computations expressed via forward()' — the distinction between model and loss is just what the forward pass computes (predictions vs. a scalar error term), not which base class it uses.