load_model.py

import torch

from models import *

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def load_model(model_path):
    """
    Load a PyTorch model from the specified path.

    Args:
    model_path (str): Path to the model file.

    Returns:
    torch.nn.Module: The loaded PyTorch model.
    """
    model = torch.load(model_path, weights_only=False)
    model.to(device)
    model.eval()  # Set the model to evaluation mode
    return model

# Example usage
if __name__ == "__main__":
    model_path = "models/nwp_model.pth"
    model = load_model(model_path)
    print("Model loaded successfully.")