Update app.py

app.py

@@ -19,7 +19,7 @@ model.eval()
 
 def predict(input_image):
     with torch.no_grad():
-        output = model(input_image.unsqueeze(0))
+        output = model(input_image.unsqueeze(0).to(model.device))
     output = output.squeeze(0).cpu().float()
     return output[0]