change pooling to mean pooling in latent space

vae.py

@@ -20,9 +20,7 @@ def __init__(self, config):
         self.activation = nn.Tanh()
 
     def forward(self, hidden_states):
-        # We "pool" the model by simply taking the hidden state corresponding
-        # to the first token.
-        first_token_tensor = hidden_states[:, 0]
+        first_token_tensor = hidden_states.mean(-1)
         pooled_output = self.dense(first_token_tensor)
         pooled_output = self.activation(pooled_output)
         return pooled_output