fix single gpu training

CLARIN-PL · Nov 30, 2023 · 90b406e · 90b406e
1 parent 9da5d26
commit 90b406e
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/embeddings/model/lightning_module/lightning_module.py b/embeddings/model/lightning_module/lightning_module.py
@@ -90,10 +90,11 @@ def predict(
 
         if return_predictions:
             assert predictions is not None
-            logits, preds = zip(*predictions)
+            logits, preds, labels = zip(*predictions)
             probabilities = softmax(torch.cat(logits), dim=1)
             preds = torch.cat(preds)
-            labels = torch.cat([x["labels"] for x in dataloader])
+            labels = torch.cat(labels)
+            # labels = torch.cat([x["labels"] for x in dataloader])
         else:
             files = sorted(os.listdir(predpath))
             all_preds = []