From 1e6683a6aa362c0cc6cf89621ce3bc064d7d7170 Mon Sep 17 00:00:00 2001
From: Andrzej Preibisz <andrzejjpreibisz@gmail.com>
Date: Sun, 12 Feb 2023 19:16:19 +0100
Subject: [PATCH] GPT2 head WIP

---
 gpt2.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/gpt2.py b/gpt2.py
index 61e2828..1aa2f73 100644
--- a/gpt2.py
+++ b/gpt2.py
@@ -84,10 +84,10 @@ class GPT2ClassificationHeadCustom(nn.Module):
         x = torch.relu(x)
         x = self.dropout(x)
         
+        x = torch.flatten(x)
         x = torch.relu(x)
-        x = self.dense_2(x)
         x = self.dropout(x)
-
+	
         x = self.out_proj(x)
         return x