Update model.py

model.py

@@ -84,5 +84,5 @@ class TransformerSentimentModel(nn.Module):
         for block in self.blocks:
             x = block(x)
 
-        x = torch.mean(x, dim=1)
+        x = x[:, 0, :]
         return self.fc(self.dropout(x))