#5 BUGFIX: missing first word id

namiyousef · Mar 6, 2022 · dd9d903 · dd9d903
1 parent 034c9fc
commit dd9d903
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/argminer/data.py b/argminer/data.py
@@ -42,7 +42,7 @@ def __getitem__(self, index):
 
         word_ids = inputs.word_ids()
         word_id_mask = [bool(word_id) for word_id in word_ids]
-        word_ids = [word_id for word_id in word_ids if word_id]
+        word_ids = [word_id for word_id in word_ids if word_id is not None]
 
         inputs = {
             key: torch.as_tensor(val, dtype=torch.long) for key, val in inputs.items()