Update app.py

app.py

@@ -40,8 +40,7 @@ def load_model(repo_id):
     # Initialize the ResNet-18 architecture
     model = torchvision.models.resnet18(pretrained=False)
     num_ftrs = model.fc.in_features
-    model.fc = nn.Linear(num_ftrs, 100)  # Adjust for your task (e.g., 128 classes)
-    # TODO: check if this number^^ corresponds to the number of classes
+    model.fc = nn.Linear(num_ftrs, 100)  # Adjust for your number of classes
 
     # Load the model weights
     state_dict = torch.load(weights_path, map_location=torch.device("cpu"))