Update app.py

app.py

@@ -1,10 +1,16 @@
 import gradio as gr
 import torch
+from torch import nn
 from torchvision import transforms
+from PIL import Image
+!wget https://huggingface.co/antonovmaxim/aiornot-kodIIm-14/resolve/main/model.pth -O model.pth
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 model = torch.hub.load("pytorch/vision", "resnet101", pretrained=False)
 model.fc = nn.Sequential(nn.Linear(2048, 500), nn.ReLU(), nn.Linear(500, 2), nn.Softmax(1))
-state_dict = torch.load('model.pth')
+state_dict = torch.load('model.pth', map_location=torch.device('cpu'))
 model.load_state_dict(state_dict)
+model.to(device)
+model.eval()
 
 transform = transforms.Compose([
     transforms.RandomHorizontalFlip(p=0.5),
@@ -15,7 +21,7 @@ transform = transforms.Compose([
 ])
 
 def classify(input_img):
-  return model(transform(input_img))[0][1].item()
+  return model(transform(Image.fromarray(input_img)).to(device).unsqueeze(0))[0][0].item()
 def img_classify(input_img):
     s = "Вероятность того, что изображение сгенерировано нейросетью равна: " + str(classify(input_img))
     return s