Update app.py

app.py

@@ -40,8 +40,12 @@ def inference(img, task, run_on):
         height = int(scale*height)
         img = img.resize((width,height), Image.ANTIALIAS)
       
-
-    model = torch.jit.load('deshabby.pt')    
+    #parameters = {'inp_channels':3, 'out_channels':3, 'dim':48, 'num_blocks':[4,6,6,8], 'num_refinement_blocks':4, 'heads':[1,2,4,8], 'ffn_expansion_factor':2.66, 'bias':False, 'LayerNorm_type':'WithBias', 'dual_pixel_task':False}
+    checkpoint = torch.load('deshabby.pt')
+    model.load_state_dict(checkpoint['params'])
+    #model.eval()
+    #model = torch.jit.load('deshabby.pt') 
+    
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')  
     model = model.to(device)
     model.eval()