Update app.py

app.py

@@ -1,16 +1,21 @@
 import torch
 from PIL import Image
-from torchvision import transforms
+from torchvision import datasets, models, transforms
 import gradio as gr
 import os
 
 
 os.system("wget https://github.com/liuxiaoyuyuyu/vanGogh-and-Other-Artist/blob/main/artist_classes.txt")
+os.system("wget https://github.com/liuxiaoyuyuyu/vanGogh-and-Other-Artist/blob/main/model_weights_mobilenet_v2_valp1trainp2.pth")
 
-model = torch.hub.load('pytorch/vision:v0.9.0', 'mobilenet_v2', pretrained=True)
+#model = torch.hub.load('pytorch/vision:v0.9.0', 'mobilenet_v2', pretrained=False)
 #checkpoint = 'https://github.com/liuxiaoyuyuyu/vanGogh-and-Other-Artist/blob/main/model_weights_mobilenet_v2_valp1trainp2.pth'
 #model.load_state_dict(torch.hub.load_state_dict_from_url(checkpoint, progress=False))
-model.eval()
+model = models.mobilenet_v2()
+num_ftrs = model.classifier[1].in_features
+model.classifier[1] = nn.Linear(num_ftrs, 6)
+model = model.to(device)
+model.load_state_dict(torch.load('model_weights_mobilenet_v2_valp1trainp2.pth'))
 
 #torch.hub.download_url_to_file("https://github.com/pytorch/hub/raw/master/images/dog.jpg", "dog.jpg")