Update app.py

app.py

@@ -138,7 +138,10 @@ class ECCVGenerator(BaseColor):
         conv8_3 = self.model8(conv7_3)
         out_reg = self.model_out(self.softmax(conv8_3))
 
-        return self.unnormalize_ab(self.upsample4(out_reg))
+        x= self.unnormalize_ab(self.upsample4(out_reg))
+        zeros = torch.zeros_like(x[:, :1, :, :])
+        x = torch.cat([x, zeros], dim=1)  # concatenate the tensor of zeros with the input tensor along the channel dimension
+        return x
 
 
 # model_net = torch.load(f'weights/colorizer.pt')