Update app.py

app.py

@@ -78,7 +78,7 @@ def inference_mask1(
     }
     nshot = support_masks.size(1)
     pred_mask, simi, simi_map = model.predict_mask_nshot(batch, nshot=nshot)
-    pred_mask = pred.detach().cpu().numpy()
+    pred_mask = pred_mask.detach().cpu().numpy()[0]
     output_img = query_img_np * 0.5 + 0.5 * np.array([1, 0, 0]) * np.expand_dims(pred_mask, axis=0)
     output_img = (output_img * 255).astype(np.uint8)
     return [output_img]