Update app.py

app.py

@@ -84,9 +84,9 @@ def predict(net,  inputs_val, shapes_val, hypar, device):
   
     inputs_val_v = Variable(inputs_val, requires_grad=False).to(device) # wrap inputs in Variable
    
-    ds_val = net(inputs_val_v)[1] # list of 6 results
+    ds_val = net(inputs_val_v)[0] # list of 6 results
 
-    pred_val = ds_val[0][0,:,:,:] # B x 1 x H x W    # we want the first one which is the most accurate prediction
+    pred_val = ds_val[1][0,:,:,:] # B x 1 x H x W    # we want the first one which is the most accurate prediction
 
     ## recover the prediction spatial size to the orignal image size
     pred_val = torch.squeeze(F.upsample(torch.unsqueeze(pred_val,0),(shapes_val[0][0],shapes_val[0][1]),mode='bilinear'))