fix?

2023-02-13 00:02:31 +01:00 · 2023-02-13 00:02:31 +01:00 · e4cdb95c3d
commit e4cdb95c3d
parent b4227f7732
1 changed files with 1 additions and 1 deletions
--- a/gpt2.py
+++ b/gpt2.py
@ -68,7 +68,7 @@ class GPT2ClassificationHeadCustom(nn.Module):
    def forward(self, x, **kwargs):
        if 'hidden_states' in kwargs and kwargs['hidden_states'] is not None:
            # Get last 5 hidden states from the end
-            hidden = torch.cat(kwargs['hidden_states'][-5:], dim=2)
+            hidden = torch.cat(kwargs['hidden_states'][-5:])
        else:
            hidden = torch.zeros(x.size(), dtype=x.dtype, device=x.device)