Patched code sample

import torch
import os

def load_model_safely(model_path: str):
    """
    Loads model weights from a file, demonstrating the fix for
    insecure deserialization.

    The vulnerability existed when torch.load was called without the
    weights_only=True parameter, allowing a malicious file to execute
    arbitrary code.
    """
    # Vulnerable call (for demonstration purposes):
    # state_dict = torch.load(model_path)

    # Fixed call:
    # By setting weights_only=True, torch.load will safely only load tensors
    # and raise an error if the file contains any other Python objects,
    # mitigating the risk of arbitrary code execution.
    if not os.path.exists(model_path):
        raise FileNotFoundError(f"Model file not found at {model_path}")
        
    state_dict = torch.load(model_path, weights_only=True)

    return state_dict

# Example usage:
# This part would typically be part of the Ludwig model server startup logic.
# To run this example, you would need a valid PyTorch model weights file (.pt or .pth).
#
# model_file = "path/to/your/model.pt"
# try:
#     print(f"Safely loading model weights from {model_file}...")
#     weights = load_model_safely(model_file)
#     print("Model weights loaded successfully.")
#     # The server would then proceed to load these weights into a model.
# except (FileNotFoundError, RuntimeError) as e:
#     print(f"Error loading model: {e}")