The matched-shape branch computed torch.linalg.norm(tensor_a, dim=1)
for every non-4-D tensor, including 1-D weights (e.g. x_pad_token,
norm .weight). A 1-D tensor has no dim=1, so this raised
'Dimension out of range ... but got 1' and aborted the whole
extraction. The code already skips 1-D tensors a few lines later, so
bail out before the norm instead.
Fixes#15