Update app.py

app.py

@@ -59,7 +59,7 @@ def predict(img):
     
     outputs = torch.argmax(outputs,1)
 
-    mask = np.array(outputs.cpu())
+    mask = np.array(outputs)
     mask[mask==1]=150
     mask[mask==3]=76 #pole # y no 74
 # mask[mask==5]=74 #pole