Update app.py

app.py

@@ -32,13 +32,13 @@ class SegModel(nn.Module):
         super().__init__()
         segformer = SegformerForSemanticSegmentation.from_pretrained(
             "nvidia/mit-b1",
-            num_labels=len(classes),
+            num_labels=len(target_list),
             id2label=id2label, 
             label2id=label2id
-        )
+        ).to(device)
 
     def forward(self, x):
-        return self.model(x).logits
+        return self.upsample(self.segformer(x).logits)
  
 model = SegModel(segformer).to(device)