SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

add easy tensor shape match test

Browse Source
This commit is contained in:
patrickvonplaten
2019-12-25 16:34:28 +01:00
parent 90cda45e9e
commit 9398058e19
1 changed files with 4 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
src/transformers/modeling_utils.py
View File

@@ -923,7 +923,10 @@ class PreTrainedModel(nn.Module):
for layer_past in past: for layer_past in past:
# copy the relevant beam idx past to past # copy the relevant beam idx past to past
reordered_layer_past = [layer_past[:, i].unsqueeze(1).clone().detach() for i in beam_idx] reordered_layer_past = [layer_past[:, i].unsqueeze(1).clone().detach() for i in beam_idx]
reordered_past.append(torch.cat(reordered_layer_past, dim=1)) reordered_layer_past = torch.cat(reordered_layer_past, dim=1)
# check that shape matches
assert reordered_layer_past.shape == layer_past.shape
reordered_past.append(reordered_layer_past)
past = tuple(reordered_past) past = tuple(reordered_past)
# update current length # update current length