How to save and load a PyTorch model

PyTorch model files move more safely between training and inference code when they store a model state_dict instead of a pickled module object. A state dictionary contains the learned tensors by layer name, so the loading code can rebuild the model class and apply only the saved weights.

The model class still has to exist in the loading code with the same layer names and tensor shapes that produced the file. torch.save(model.state_dict(), “weather-score-state-dict.pth”) writes the parameters, and load_state_dict() copies those tensors into a fresh module instance.

Use torch.load(…, weights_only=True) for state dictionaries and keep map_location=“cpu” when a file may have been saved from another device. After loading, switch the module to evaluation mode before inference so dropout and batch normalization layers use inference behavior.

1. Create a smoke script that saves a state dictionary and loads it into a new model instance.

   save_load_model.py

   import torch
   from torch import nn
    
    
   torch.manual_seed(11)
   model_path = "weather-score-state-dict.pth"
    
    
   class WeatherScoreModel(nn.Module):
       def __init__(self):
           super().__init__()
           self.net = nn.Sequential(
               nn.Linear(3, 5),
               nn.ReLU(),
               nn.Linear(5, 1),
           )
    
       def forward(self, inputs):
           return self.net(inputs)
    
    
   def build_model():
       model = WeatherScoreModel()
       model.eval()
       return model
    
    
   sample = torch.tensor([[0.6, 0.2, 0.9]], dtype=torch.float32)
    
   trained_model = build_model()
   with torch.inference_mode():
       expected_output = trained_model(sample)
    
   torch.save(trained_model.state_dict(), model_path)
   print(f"saved={model_path}")
    
   loaded_model = build_model()
   state_dict = torch.load(model_path, map_location="cpu", weights_only=True)
   load_result = loaded_model.load_state_dict(state_dict)
   loaded_model.eval()
    
   with torch.inference_mode():
       loaded_output = loaded_model(sample)
    
   print(f"keys={len(state_dict)}")
   print(f"load_result={load_result}")
   print(f"outputs_match={torch.allclose(expected_output, loaded_output)}")
   print(f"output_shape={tuple(loaded_output.shape)}")
   print(f"prediction={loaded_output.item():.4f}")

   The smoke model uses deterministic initial weights so the save/load check is repeatable. Replace WeatherScoreModel with the class that created your trained state dictionary.

2. Run the smoke script.

   $ python save_load_model.py
   saved=weather-score-state-dict.pth
   keys=4
   load_result=<All keys matched successfully>
   outputs_match=True
   output_shape=(1, 1)
   prediction=-0.0112

   outputs_match=True confirms that the loaded model returns the same inference result as the original model for the sample input.

3. Save only the trained model parameters in the real training or fine-tuning code.

   torch.save(model.state_dict(), "model-state-dict.pth")

   Use .pth or .pt consistently with your project conventions. A state dictionary file is less tied to source-code paths than a pickled full module.

4. Recreate the model class before loading the state dictionary.

   model = WeatherScoreModel()
   state_dict = torch.load(
       "model-state-dict.pth",
       map_location="cpu",
       weights_only=True,
   )
   model.load_state_dict(state_dict)
   model.eval()

   Load model files only from sources you trust. weights_only=True narrows unpickling to tensors, primitive values, dictionaries, and allowlisted safe globals, but it is not a general malware or denial-of-service boundary.

5. Run inference under torch.inference_mode() after the load.

   sample = torch.tensor([[0.6, 0.2, 0.9]], dtype=torch.float32)
    
   with torch.inference_mode():
       prediction = model(sample)
    
   print(prediction.shape)

   Related: How to run inference in PyTorch

6. Remove the smoke-test files after confirming the pattern.

   $ rm save_load_model.py weather-score-state-dict.pth