mirror of
https://github.com/facebookresearch/fairseq.git
synced 2024-10-26 17:32:57 +03:00
Fix indexing in TokenBlockDataset
Summary: Pull Request resolved: https://github.com/pytorch/fairseq/pull/719 Differential Revision: D15258483 Pulled By: myleott fbshipit-source-id: dd00daa6f1c87264c1196a77dfffc8c876ebde7f
This commit is contained in:
parent
0cb45bcb12
commit
eddcdf08e1
@ -70,7 +70,7 @@ class TokenBlockDataset(FairseqDataset):
|
||||
if not torch.is_tensor(sizes):
|
||||
sizes = torch.tensor(sizes)
|
||||
cumsum = torch.cumsum(sizes, dim=0)
|
||||
self.slice_indices[0, 1] = sizes[0]
|
||||
self.slice_indices[0] = [0, sizes[0]]
|
||||
self.slice_indices[1:] = cumsum.unfold(0, 2, 1)
|
||||
else:
|
||||
raise ValueError('Invalid break_mode: ' + break_mode)
|
||||
|
Loading…
Reference in New Issue
Block a user