refactor: add KeyError handling for DataBatch

vejvarm · Nov 27, 2023 · 8933e58 · 8933e58
1 parent 9cc69fc
commit 8933e58
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/dataset.py b/dataset.py
@@ -66,7 +66,7 @@ def __init__(self, batch: list[list[any]], vocabs: dict, device: str):
                 entity_pointer.append(self._tensor([vocabs[ENTITY].stoi[s] for s in sample[7]]))
                 KeyError: 'Q4622539'
                 """
-                entity_pointer.append(self._tensor([vocabs['NA'].stoi[s] for s in sample[7]]))
+                entity_pointer.append(self._tensor([vocabs['ENTITY'].stoi['NA']]))
         self.id = self._tensor(id).to(device)
         self.input = pad_sequence(inp,
                                   padding_value=vocabs[INPUT].stoi[PAD_TOKEN],