Update app.py

app.py

@@ -15,7 +15,10 @@ def inference(repo_id, model_name, img):
     with torch.inference_mode():
         argmax, softmax, stdeviation = model(inp)
     
-    return argmax*255, stdeviation
+    pred = argmax.cpu().numpy()
+    std = stdeviation.cpu().numpy()
+    
+    return pred*255, std
  
     
 title="deepflash2"