Minor update - needed for this to run on an older version of torch

stanfordnlp · Oct 30, 2023 · 90fee9a · 90fee9a
1 parent 9b79719
commit 90fee9a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/stanza/models/coref/model.py b/stanza/models/coref/model.py
@@ -415,7 +415,7 @@ def _build_optimizers(self):
     def _clusterize(self, doc: Doc, scores: torch.Tensor, top_indices: torch.Tensor):
         antecedents = scores.argmax(dim=1) - 1
         not_dummy = antecedents >= 0
-        coref_span_heads = torch.arange(0, len(scores))[not_dummy]
+        coref_span_heads = torch.arange(0, len(scores), device=not_dummy.device)[not_dummy]
         antecedents = top_indices[coref_span_heads, antecedents[not_dummy]]
 
         nodes = [GraphNode(i) for i in range(len(doc["cased_words"]))]