Update app.py

app.py

@@ -1,12 +1,14 @@
 import gradio as gr
 from torch.nn import functional as F
+import torch
 
 from model_loader import ModelType, type_to_transforms, type_to_loaded_model
 
 def get_y(model_type, model, image):
-    if model_type == ModelType.SYNTHETIC_DETECTOR_V2:
-        return model.forward(image.unsqueeze(0).to("cpu"))
-    return model.forward(image[None, ...])
+    with torch.no_grad():
+        if model_type == ModelType.SYNTHETIC_DETECTOR_V2:
+            return model.forward(image.unsqueeze(0).to("cpu"))
+        return model.forward(image[None, ...])
 
 def predict(raw_image, model_name):
     if model_name not in ModelType.get_list():
@@ -17,10 +19,11 @@ def predict(raw_image, model_name):
     tfm = type_to_transforms[model_type]
     image = tfm(raw_image)
     y = get_y(model_type, model, image)
-    y_1 = F.softmax(y, dim=1)[:, 1].cpu().detach().numpy()
-    y_2 = F.softmax(y, dim=1)[:, 0].cpu().detach().numpy()
-    return {'created by AI': y_1.tolist(),
-            'created by human': y_2.tolist()}
+
+    y_2, y_1 = F.softmax(y, dim=1).cpu().numpy()[0]
+    
+    return {'created by AI': float(y_2),
+            'created by human': float(y_1)}
 
 general_examples = [
     ["images/general/img_1.jpg"],