Update amplify to extract last embedding after last layer norm.

amplify.py

@@ -276,7 +276,10 @@ class AMPLIFY(AMPLIFYPreTrainedModel):
                 attentions.append(attn)
 
         # Classification head with layer norm
-        logits = self.decoder(self.layer_norm_2(x) if self.config.layer_norm_before_last_layer else x)
+        x_normalized = self.layer_norm_2(x) if self.config.layer_norm_before_last_layer else x
+        if output_hidden_states:
+            hidden_states[-1]=x_normalized
+        logits = self.decoder(x_normalized)
 
         # Return logits or the output of the last hidden layer
         return MaskedLMOutput(logits=logits, hidden_states=hidden_states, attentions=attentions)