Update app.py

app.py

@@ -20,7 +20,7 @@ LESION_COLORS = {
 
 # ====================== UNET ARCHITECTURE ======================
 class UNet(nn.Module):
-    def __init__(self, input_channels=3, num_classes=5):
+    def __init__(self, input_channels=3, num_classes=3):
         super(UNet, self).__init__()
         
         def conv_block(in_channels, out_channels):