Fix VarBuilder handling in GTE e.g. gte-multilingual-reranker-base (#538)

Co-authored-by: Alvaro Bartolome <36760800+alvarobartt@users.noreply.github.com>

Author: Narsil

backends/candle/src/models/gte.rs

@@ -408,6 +408,42 @@ impl GTEModel {
             ModelType::Embedding(pool) => (pool, None),
         };
 
+        let (word_embeddings, token_type_embeddings, encoder, embeddings_norm) =
+            Self::inner_load(vb.pp("new"), config)
+                .or_else(|_| Self::inner_load(vb.clone(), config))?;
+
+        let rotary_dim = encoder.layers[0].attention.attention_head_size;
+        let inv_freqs = get_inv_freqs(
+            rotary_dim,
+            config.rope_theta,
+            vb.device(),
+            config.rope_scaling.as_ref(),
+        )?;
+
+        let rotary_cache =
+            get_cos_sin(config.max_position_embeddings, &inv_freqs, vb.dtype(), true)?;
+
+        Ok(Self {
+            word_embeddings,
+            token_type_embeddings,
+            encoder,
+            embeddings_norm,
+            rotary_cache,
+            classifier,
+            pool,
+            num_attention_heads: config.num_attention_heads,
+            device: vb.device().clone(),
+            dtype: vb.dtype(),
+            span: tracing::span!(tracing::Level::TRACE, "model"),
+            rotary_dim,
+        })
+    }
+
+    fn inner_load(
+        vb: VarBuilder,
+        config: &GTEConfig,
+    ) -> Result<(Embedding, Option<Embedding>, GTEEncoder, LayerNorm)> {
+
         let word_embeddings = Embedding::new(
             vb.pp("embeddings.word_embeddings")
                 .get((config.vocab_size, config.hidden_size), "weight")?,
@@ -431,32 +467,12 @@ impl GTEModel {
             config.hidden_size,
             config.layer_norm_eps,
         )?;
-
-        let rotary_dim = encoder.layers[0].attention.attention_head_size;
-        let inv_freqs = get_inv_freqs(
-            rotary_dim,
-            config.rope_theta,
-            vb.device(),
-            config.rope_scaling.as_ref(),
-        )?;
-
-        let rotary_cache =
-            get_cos_sin(config.max_position_embeddings, &inv_freqs, vb.dtype(), true)?;
-
-        Ok(Self {
+        Ok((
             word_embeddings,
             token_type_embeddings,
             encoder,
             embeddings_norm,
-            rotary_cache,
-            classifier,
-            pool,
-            num_attention_heads: config.num_attention_heads,
-            device: vb.device().clone(),
-            dtype: vb.dtype(),
-            span: tracing::span!(tracing::Level::TRACE, "model"),
-            rotary_dim,
-        })
+        ))
     }
 
     pub fn forward(&self, batch: Batch) -> Result<(Option<Tensor>, Option<Tensor>)> {