Update app.py

app.py

@@ -8,9 +8,13 @@ inference = torch.load('fine_tune_resnet.pth')
 inference.eval()
 
 def classifier(image):
-    output = inference(test_transform(Image.open('/content/1000-ml-plastic-water-bottle-500x500.webp')).unsqueeze(0))
-    _, prediction = torch.max(output,1)
-    confidence = round(torch.softmax(output,1).max().item(),4)*100
+    if torch.cuda.is_available():
+        inference.cpu()
+
+    with torch.no_grad():
+        output = inference(test_transform(Image.open('/content/1000-ml-plastic-water-bottle-500x500.webp')).unsqueeze(0))
+        _, prediction = torch.max(output,1)
+        confidence = round(torch.softmax(output,1).max().item(),4)*100
 
     return f'{label_dict[prediction.item()]} (Confidence: {confidence}%)'