implement model

app.py

@@ -44,7 +44,7 @@ def load_model(repo_id):
     # Initialize the ResNet-18 architecture
     model = torchvision.models.resnet18(pretrained=True) # TODO: does it matter if this is set to true or false?
     num_ftrs = model.fc.in_features
-    model.fc = nn.Linear(num_ftrs, 100)  # Adjust for your task (e.g., 128 classes)
+    model.fc = nn.Linear(num_ftrs, 128)  # Adjust for your task (e.g., 128 classes)
     # TODO: check if this number^^ corresponds to the number of classes
 
     # Load the model weights
@@ -78,10 +78,13 @@ def process_image_str(groupImageFilePath: str):
         # do AI stuff here
         with torch.no_grad():
             outputs_t = model(intputTensor)
-            _, pred_t = torch.max(outputs_t, dim=1)
+            print(outputs_t)
+            temp, pred_t = torch.max(outputs_t, dim=1)
+            print(temp)
             outputLabels.append(pred_t.item())
 
     #return gr.Image(image)
+    print(outputLabels)
     return outputLabels.pop(0)