Update app.py

app.py

@@ -72,7 +72,7 @@ class FaceEmbeddingModel(torch.nn.Module):
 model_pretrained = torch.load(model_path, map_location=device, weights_only=False)
 
 # Define the ML model - Evaluation function
-def prod_function(transformer_model, prod_dl, prod_data):
+def prod_function(transformer_model, prod_dl, webcam_dl):
     # Initialize accelerator
     accelerator = Accelerator()
 
@@ -88,22 +88,23 @@ def prod_function(transformer_model, prod_dl, prod_data):
     set_seed(42)
 
     # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the prepare method.
-    accelerated_model, acclerated_prod_dl, acclerated_prod_data = accelerator.prepare(transformer_model, prod_dl, prod_data)
+    accelerated_model, acclerated_prod_dl, acclerated_webcam_dl = accelerator.prepare(transformer_model, prod_dl, webcam_dl)
 
     # Evaluate at the end of the epoch
     accelerated_model.eval()
 
     # Find Embedding of the image to be evaluated
-    with torch.no_grad():
-        img_prod = acclerated_prod_data['pixel_values']
-        emb_prod = accelerated_model(img_prod)
+    for batch in acclerated_webcam_dl:
+        with torch.no_grad():
+            #img_prod = acclerated_prod_data['pixel_values']
+            emb_prod = accelerated_model(batch)
 
     prod_preds = []
 
     for batch in acclerated_prod_dl:
-        img = batch['pixel_values']
+        #img = batch['pixel_values']
         with torch.no_grad():
-            emb = accelerated_model(img)
+            emb = accelerated_model(batch)
         distance = F.pairwise_distance(emb, emb_prod)
 
         prod_preds.append(distance)