import torch
import torch.nn as nn
from torchvision import models


def load_model(pretrained_weights_path):
    # Initialize Face-Rego
    net = models.resnet18(pretrained=False)
    num_ftrs = net.fc.in_features
    net.fc = nn.Linear(num_ftrs, 100)  # Match your fine-tuned setup

    # Load weights
    state_dict = torch.load(pretrained_weights_path, map_location=torch.device('cpu'))
    net.load_state_dict(state_dict)
    net.eval()  # Set to evaluation mode
    return net