Update app.py

app.py

@@ -43,7 +43,7 @@ dog_breeds = ["Afghan_Hound", "African_Hunting_Dog", "Airedale", "American_Staff
               "Miniature_Poodle", "Miniature_Schnauzer", "Otterhound", "Papillon", "Pug", "Redbone", 
               "Schipperke", "Silky_Terrier", "Soft-Coated_Wheaten_Terrier", "Standard_Poodle", 
               "Standard_Schnauzer", "Toy_Poodle", "Toy_Terrier", "Vizsla", "Whippet", 
-              "Wire-Haired_Fox_Terrier"]
+              "Wire-Haired_Fox_Terrier","Havanese", "Dachshund", "Shiba_Inu", "Bichon_Frise"]
 
 class MultiHeadAttention(nn.Module):
 
@@ -100,11 +100,11 @@ class BaseModel(nn.Module):
         return logits, attended_features
 
 
-num_classes = 120
+num_classes = 124
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 model = BaseModel(num_classes=num_classes, device=device)
 
-checkpoint = torch.load('best_model_81_dog.pth', map_location=torch.device('cpu'))
+checkpoint = torch.load('[124_82.30]_best_model_dog.pth', map_location=torch.device('cpu'))
 model.load_state_dict(checkpoint['model_state_dict'])
 
 # evaluation mode