Commit 51eee004 authored by Gaëtan Caillaut's avatar Gaëtan Caillaut
Browse files

fix out of bounds error

parent b4fc174c
......@@ -77,7 +77,7 @@ class TrainData:
attention_mask_tensor[i, :] = torch.tensor(
encoded.attention_mask[:self.max_seq_size])
for j, wid in enumerate(encoded.word_ids):
for j, wid in enumerate(encoded.word_ids[:self.max_seq_size]):
if wid is None:
break
word_ids_tensor[i, j] = wid
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment