refactor: normalization after truncation

Signed-off-by: jupyterjazz <saba.sturua@jina.ai>

modeling_xlm_roberta.py

@@ -588,12 +588,7 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
                     embeddings = self.mean_pooling(
                         token_embs, encoded_input["attention_mask"]
                     )
-
-                if normalize_embeddings:
-                    embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
-
-                if convert_to_numpy:
-                    embeddings = embeddings.cpu()
+            
             all_embeddings.extend(embeddings)
 
         all_embeddings = [all_embeddings[idx] for idx in inverse_permutation]
@@ -601,11 +596,14 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
         truncate_dim = truncate_dim or self.config.truncate_dim
         if truncate_dim:
             all_embeddings = self.truncate_embeddings(all_embeddings, truncate_dim)
+        
+        if normalize_embeddings:
+            all_embeddings = [torch.nn.functional.normalize(embedding, p=2, dim=0) for embedding in all_embeddings]
 
         if convert_to_tensor:
             all_embeddings = torch.stack(all_embeddings)
         elif convert_to_numpy:
-            all_embeddings = np.asarray([emb.numpy() for emb in all_embeddings])
+            all_embeddings = np.asarray([emb.cpu().numpy() for emb in all_embeddings])
 
         if input_was_string:
             all_embeddings = all_embeddings[0]