Upload model.py with huggingface_hub

model.py

@@ -0,0 +1,16 @@
+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, 128)  # 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