try fix preds gpu

CLARIN-PL · Nov 30, 2023 · 73dbbff · 73dbbff
1 parent e6dd22f
commit 73dbbff
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/embeddings/model/lightning_module/lightning_module.py b/embeddings/model/lightning_module/lightning_module.py
@@ -106,7 +106,7 @@ def predict(
                     with open(os.path.join(predpath, file), "rb") as f:
                         batch_indices = pickle.load(f)
                         all_batch_indices.append(list(flatten(batch_indices)))
-            all_batch_indices = torch.Tensor([y for x in all_batch_indices for y in x])
+            all_batch_indices = torch.Tensor([y for x in all_batch_indices for y in x]).long()
             probabilities = softmax(torch.cat(all_logits), dim=1)[all_batch_indices]
             preds = torch.cat(all_preds)[all_batch_indices]