Spaces:

Alic22
/

Masterarbeit

Sleeping

Alic22 commited on Jan 5

Commit

749b7d2

verified ·

1 Parent(s): 5e89f53

Update train.py

Files changed (1) hide show

train.py CHANGED Viewed

@@ -56,7 +56,7 @@ loader = DataLoader(dataset, batch_size=2, shuffle=True)
 # ------------------
 segformer = SegformerForSemanticSegmentation.from_pretrained(
     "nvidia/mit-b1",
-    num_labels=13,
     id2label=id2label,
     label2id=label2id
 ).to(device)

 # ------------------
 segformer = SegformerForSemanticSegmentation.from_pretrained(
     "nvidia/mit-b1",
+    num_labels=len(target_list),
     id2label=id2label,
     label2id=label2id
 ).to(device)