Update model.py

model.py

@@ -84,7 +84,7 @@ class Model(LabelStudioMLBase):
                         logits = torch.exp(logits)
                         print(logits)
                         preds = logits > 0.5
-                        preds = [self.seg_id2label[i] for i, pred in enumerate(preds) if pred.item()]
+                        preds = [self.seg_id2label[i] for i, pred in enumerate(preds.squeeze().tolist()) if pred.item()]
                         preds = ["No Reportado"] if "No Reportado" in preds else preds
                         result_list.append({
                           "value": {