Table of Contents
What Changes Between Training and Inference
Inference means using a trained model to produce predictions on new data. In PyTorch, the two most important switches are setting the model to evaluation mode and turning off gradient tracking. Evaluation mode is activated with model.eval() and it changes the behavior of some layers such as dropout and batch normalization. Disabling gradients is done with torch.no_grad() for most inference code, or torch.inference_mode() for an even more optimized inference context in many cases.
Always call model.eval() before inference, and run the forward pass inside torch.no_grad() or torch.inference_mode(). Otherwise you can get incorrect predictions and unnecessary memory use.
A Minimal Inference Function
A practical pattern is to wrap preprocessing, device placement, the forward pass, and postprocessing into one function. Your preprocessing must match what you used during training, especially for feature scaling, categorical encoding, and column ordering. For structured data, mismatched feature order is one of the most common silent bugs.
A typical flow is to move inputs to the same device as the model, run the model, then convert outputs back to CPU for downstream use. If you trained a classifier, you often want probabilities or predicted classes. If you trained a regressor, you often want raw numeric predictions.
Saving and Loading for Inference, Weights-Only
For beginners, the most common and recommended approach is to save and load the model weights with a state_dict. You save after training with torch.save(model.state_dict(), path) and later recreate the model object with the same architecture code, load weights with model.load_state_dict(torch.load(path, map_location=device)), then switch to model.eval().
This workflow is simple and robust, but it requires that the code defining the model class is available at inference time, because you must instantiate the architecture before loading weights.
A state_dict only stores parameters and buffers, not the model class definition. You must rebuild the same architecture before calling load_state_dict.
Saving and Loading Full Checkpoints, When You Need to Resume Training
Sometimes you want more than inference. You want to resume training exactly, including optimizer state, learning rate schedule, and epoch counters. In that case, you save a checkpoint dictionary that includes model_state_dict, optimizer_state_dict, and any metadata like epoch and best validation metric.
For pure inference, a full checkpoint is usually unnecessary. It is larger and easier to misuse, for example by accidentally loading an optimizer you do not need.
Device Placement at Inference Time
Inference can run on CPU, CUDA GPU, or Apple Silicon. The key rule is that the model parameters and the input tensors must be on the same device. If you load weights on CPU but your input is on GPU, or the reverse, you will get a device mismatch error. A common safe pattern is to select a device, load the weights with map_location=device, move the model to that device, then move each batch of inputs to that device right before the forward pass.
Output Postprocessing for Common Structured Data Tasks
For binary classification, models often output a single logit per example. To convert logits to probabilities, you apply the sigmoid function, $p = \sigma(z) = \frac{1}{1 + e^{-z}}$, and to convert probabilities to a hard class label you apply a threshold, often $0.5$. For multi-class classification, models often output a vector of logits per example, and you apply softmax to get probabilities and argmax to get the predicted class index. For regression, you usually use the model output directly, possibly reversing any target scaling you applied during training.
Use the right postprocessing for the loss you trained with. If you trained with BCEWithLogitsLoss, your model outputs logits, not probabilities, so apply sigmoid at inference time when you need probabilities.
Batch Inference With a DataLoader
For structured datasets, batch inference is typically faster and simpler than predicting one row at a time. You can reuse a Dataset and DataLoader to iterate over batches, run forward passes, and concatenate predictions. During inference you usually disable shuffling so predictions align with input order, unless you explicitly track indices.
Keeping batch sizes moderate helps memory usage, especially on GPU. If you see out of memory errors, reduce batch size first.
Exporting a Model, What “Export” Means Here
Saving a state_dict is not the same as exporting a model to a different runtime. Export typically means turning your model into a representation that can run outside eager PyTorch Python execution. In PyTorch courses, the two common paths you will encounter later are TorchScript and ONNX. At a beginner level, it is enough to understand why export exists: it can reduce Python dependency, improve portability, and help integrate with production systems.
For structured data feedforward networks, the main practical requirement for export readiness is that your forward pass is deterministic and uses standard tensor operations. Complicated Python side effects inside forward, or data dependent control flow that cannot be traced, can make export harder.
A Simple Inference Checklist for Structured Data
Before you trust your predictions, confirm that the feature columns are in the same order as training, that all preprocessing steps are identical, and that categorical mappings match exactly. Then verify that model.eval() is set, that you are using no_grad or inference_mode, and that device placement is consistent. Finally, sanity check by running inference on a few examples where you know the expected behavior, such as training samples where the model should perform reasonably well.