Update app.py

app.py

@@ -43,7 +43,7 @@ def greet(image):
 
     print(image.shape)
     print(mask_pred.shape)
-    return mask_pred[0].numpy()
+    return mask_pred[0, 0, :, :].numpy()
     
 
 demo = gr.Interface(