Update app.py

app.py

@@ -226,7 +226,7 @@ def colorize2(img: Image.Image, model_option: str):
         model_int = 0
     input = torch.from_numpy(np.asarray(img))[None,None].float().to(device) / 255.0 * 2 - 1
     with torch.inference_mode():
-        out2 = model_net(input)
+        out2 = model_net(input).squeeze()
         out2 = sample[0].cpu().numpy().transpose([1,2,0])
         out2 = np.uint8(((out + 1) / 2 * 255).clip(0,255))