Creating new bool tensor straight away instead of regular clone

centre-for-humanities-computing · Oct 31, 2023 · fbe8bcf · fbe8bcf
1 parent 3f474f8
commit fbe8bcf
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/paper/src/ents_heads_extraction.py b/paper/src/ents_heads_extraction.py
@@ -1,7 +1,6 @@
 """Pipeline for headwords/entities extractions and frequency count."""
 import spacy
 
-from relationextraction import SpacyRelationExtractor  # noqa
 from conspiracies.docprocessing.headwordextraction import contains_ents
 
 

diff --git a/src/conspiracies/docprocessing/relationextraction/multi2oie/other/bio.py b/src/conspiracies/docprocessing/relationextraction/multi2oie/other/bio.py
@@ -32,10 +32,10 @@ def get_pred_mask(tensor):
         where B is the batch size, L is the sequence length.
     :return: masked binary tensor with the same shape.
     """
-    res = tensor.clone()
-    res[tensor == pred_tag2idx["O"]] = 1
-    res[tensor != pred_tag2idx["O"]] = 0
-    return res.bool()
+    res = tensor.bool()
+    res[tensor == pred_tag2idx["O"]] = True
+    res[tensor != pred_tag2idx["O"]] = False
+    return res
 
 
 def filter_pred_tags(pred_tags, tokens):