How to run inference in PyTorch

A trained PyTorch model becomes useful outside training when it can turn a prepared tensor batch into scores, probabilities, or labels. Inference is the handoff from learned weights to application code, so the model has to use evaluation behavior and the forward pass should avoid gradient bookkeeping.

model.eval() changes the module's training flag and switches affected layers such as Dropout and BatchNorm into evaluation behavior. torch.inference_mode() disables autograd tracking for the forward pass and removes extra overhead when the output will not feed a later backward() call.

The smoke script uses a small CPU classifier so the inference path can be verified without a trained checkpoint. In a real project, load or instantiate the same model class that produced the weights, prepare tensors with the shape and device the model expects, and treat the output shape plus predicted labels as the first proof that the handoff works.

1. Create a minimal inference smoke script.

   inference_run_demo.py

   import torch
   from torch import nn
    
    
   class TicketClassifier(nn.Module):
       def __init__(self):
           super().__init__()
           self.dropout = nn.Dropout(p=0.5)
           self.linear = nn.Linear(4, 3)
    
       def forward(self, inputs):
           return self.linear(self.dropout(inputs))
    
    
   class_names = ["standard", "priority", "urgent"]
   model = TicketClassifier()
    
   with torch.no_grad():
       model.linear.weight.copy_(
           torch.tensor(
               [
                   [0.4, -0.2, 0.1, 0.3],
                   [-0.1, 0.5, 0.2, -0.2],
                   [0.2, 0.1, -0.1, 0.6],
               ],
               dtype=torch.float32,
           )
       )
       model.linear.bias.copy_(torch.tensor([0.0, 0.1, -0.1]))
    
   features = torch.tensor(
       [
           [0.9, 0.1, 0.2, 0.3],
           [0.1, 0.8, 0.4, 0.2],
           [0.2, 0.2, 0.1, 0.9],
       ],
       dtype=torch.float32,
   )
    
   model.eval()
    
   with torch.inference_mode():
       inference_mode_active = torch.is_inference_mode_enabled()
       logits = model(features)
       probabilities = torch.softmax(logits, dim=1)
       predicted_indexes = probabilities.argmax(dim=1)
       predicted_labels = [class_names[index] for index in predicted_indexes.tolist()]
    
   print(f"model training mode: {model.training}")
   print(f"dropout training mode: {model.dropout.training}")
   print(f"inference mode active: {inference_mode_active}")
   print(f"logits shape: {tuple(logits.shape)}")
   print(f"probabilities shape: {tuple(probabilities.shape)}")
   print(f"predicted labels: {', '.join(predicted_labels)}")
   print(f"logits require grad: {logits.requires_grad}")

   The Dropout layer makes the evaluation-mode check visible. Replace TicketClassifier and features with the model class and input tensor shape used by the trained model.

2. Run the inference smoke script.

   $ python inference_run_demo.py
   model training mode: False
   dropout training mode: False
   inference mode active: True
   logits shape: (3, 3)
   probabilities shape: (3, 3)
   predicted labels: standard, priority, urgent
   logits require grad: False

   model training mode: False and dropout training mode: False confirm evaluation behavior. inference mode active: True and logits require grad: False confirm that the forward pass did not track gradients.

3. Add the evaluation guard to the project inference function.

   def run_model(model, inputs):
       model.eval()
    
       with torch.inference_mode():
           logits = model(inputs)
    
       return logits

   model.eval() is sticky until model.train() is called. Switch back to training mode before another optimizer step.

4. Convert logits into class probabilities and labels.

   def classify(model, inputs, class_names):
       model.eval()
    
       with torch.inference_mode():
           logits = model(inputs)
           probabilities = torch.softmax(logits, dim=1)
           predicted_indexes = probabilities.argmax(dim=1)
    
       return [
           {
               "label": class_names[index],
               "confidence": probabilities[row, index].item(),
           }
           for row, index in enumerate(predicted_indexes.tolist())
       ]

   Use softmax(dim=1) for a batch of class logits shaped like batch_size x class_count. Regression models usually return raw numeric predictions instead of class probabilities.

5. Keep the model and input tensors on the same device.

   device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
   model = model.to(device)
   inputs = inputs.to(device)
    
   with torch.inference_mode():
       logits = model(inputs)

   Mixed devices raise runtime errors during the forward pass. Choose the device once near the application entry point and move both model parameters and input tensors before inference.

6. Remove the temporary smoke script.

   $ rm inference_run_demo.py