How to run inference in PyTorch

Application code turns learned parameters into a prediction only when the input shape, model state, and output interpretation agree. A PyTorch inference path therefore has to prepare the model for evaluation, run the forward pass without gradient bookkeeping, and translate the returned tensor into the value the application needs.

The model.eval() call changes modules such as Dropout and BatchNorm to evaluation behavior, while torch.inference_mode() disables autograd tracking and additional tensor bookkeeping for the enclosed operations. Neither call replaces the other, so both belong around a prediction-only forward pass.

The CPU classifier below uses fixed parameters and a three-row feature batch so every stage can be reproduced without a checkpoint download. Replace the classifier, state dictionary, class names, and feature preparation with the artifacts from the trained project while keeping the evaluation and inference boundary intact.

1. Create inference_run.py with the classifier structure and class index order.

   inference_run.py

   import torch
   from torch import nn
    
    
   class TicketClassifier(nn.Module):
       def __init__(self):
           super().__init__()
           self.layers = nn.Sequential(
               nn.Linear(4, 3),
               nn.Dropout(p=0.5),
           )
    
       def forward(self, inputs):
           return self.layers(inputs)
    
    
   class_names = ["standard", "priority", "urgent"]
   device = torch.device("cpu")
   model = TicketClassifier().to(device)

   The class_names sequence must match the index order used by the trained classifier. An accelerator path requires the model and inputs on the same selected project device.

2. Append the classifier parameters below the model creation in inference_run.py.

   model.load_state_dict(
       {
           "layers.0.weight": 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,
           ),
           "layers.0.bias": torch.tensor([0.0, 0.1, -0.1]),
       }
   )

   A real inference entry point normally loads a saved state_dict whose keys and tensor shapes match the model class.

3. Append the feature batch below the loaded parameters in inference_run.py.

   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,
       device=device,
   )

   Each row contains the four features expected by TicketClassifier. Project inputs need the same transforms and feature order used during training.

4. Append the evaluation forward pass and class conversion below the feature batch.

   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()]
   probability_sums = probabilities.sum(dim=1)

   dim=1 normalizes the class scores across each batch row. Regression and multilabel models require output handling that matches their training objective instead of this single-label softmax conversion.

5. Append fail-capable inference guards below the converted labels.

   if model.training:
       raise RuntimeError("The model is still in training mode")
   if logits.requires_grad:
       raise RuntimeError("The inference output still tracks gradients")
   if logits.shape != (3, 3):
       raise RuntimeError(f"Unexpected logits shape: {tuple(logits.shape)}")
   if not torch.allclose(probability_sums, torch.ones(3, device=device)):
       raise RuntimeError(f"Probabilities are not normalized: {probability_sums}")

6. Append the observed inference fields below the guards.

   print(f"model_training={model.training}")
   print(f"inference_mode_active={inference_mode_active}")
   print(f"logits_shape={tuple(logits.shape)}")
   print(f"probability_sums={probability_sums.tolist()}")
   print(f"predicted_labels={predicted_labels}")
   print(f"logits_require_grad={logits.requires_grad}")

7. Run the completed classifier from its project directory to confirm the inference result.

   $ python3 inference_run.py
   model_training=False
   inference_mode_active=True
   logits_shape=(3, 3)
   probability_sums=[1.0, 1.0, 1.0]
   predicted_labels=['standard', 'priority', 'urgent']
   logits_require_grad=False

   model_training=False confirms evaluation behavior, and logits_require_grad=False confirms that the returned tensor has no autograd history. The shape, normalized probability sums, and labels prove the complete batch-to-prediction path.