Update README.md

README.md

@@ -68,20 +68,35 @@ We present EdgeFace- a lightweight and efficient face recognition network inspir
     # load model
     model_name="edgeface_base"
     model=get_model(model_name)
-    checkpoint_path=f'checkpoints/{arch}.pt'
-    model.load_state_dict(torch.load(checkpoint_path, map_location='cpu')).eval()
+    checkpoint_path=f'checkpoints/{model_name}.pt'
+    # 1. Load the weights into the model
+    model.load_state_dict(torch.load(checkpoint_path, map_location='cpu'))
+
+    # 2. Set the model to evaluation mode on its own line
+    model.eval()
 
     transform = transforms.Compose([
-                transforms.ToTensor(),
-                transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
-                ])
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
+            ])
+
+    path = 'path to_face_image'
+    aligned = align.get_aligned_face(path) 
+    transformed_input = transform(aligned) 
+
+    # FIX: Add the batch dimension (turns [3, 112, 112] into [1, 3, 112, 112])
+    transformed_input = transformed_input.unsqueeze(0)
 
-    path = 'path_to_face_image'
-    aligned = align.get_aligned_face(path) # align face
-    transformed_input = transform(aligned) # preprocessing
+    # Move to the same device as the model (crucial if you switched to 'mps' or 'cuda')
+    device = next(model.parameters()).device
+    transformed_input = transformed_input.to(device)
 
     # extract embedding
+    with torch.no_grad(): # Use no_grad to save memory during inference
     embedding = model(transformed_input)
+
+    print(embedding.shape)
+
 ```