added

app.py

@@ -19,7 +19,7 @@ def predict(inputImage: torch.Tensor, model) -> torch.Tensor:
     with torch.no_grad():
         output = model(preprocessing(inputImage.clone()).to(device))
     output = output[0] if isinstance(output, (tuple, list)) else output
-    return output[0].argmax(dim=0, keepdim=True).cpu()  
+    return output[0].argmax(dim=0, keepdim=True).to(device)  
 
 
 # %% Gradio interface