fix: cast indices tensor to int to fix bug

flairNLP · Jan 25, 2025 · 32c875b · 32c875b
1 parent 30974f2
commit 32c875b
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/flair/nn/model.py b/flair/nn/model.py
@@ -899,7 +899,9 @@ def predict(
 
                         if has_unknown_label:
                             has_any_unknown_label = True
-                            scores = torch.index_select(scores, 0, torch.tensor(filtered_indices, device=flair.device))
+                            scores = torch.index_select(
+                                scores, 0, torch.tensor(filtered_indices, device=flair.device, dtype=torch.int32)
+                            )
 
                         gold_labels = self._prepare_label_tensor([data_points[index] for index in filtered_indices])
                         overall_loss += self._calculate_loss(scores, gold_labels)[0]