[rust] Fix bert model classifier loading (#3441)

deepjavalibrary · Aug 26, 2024 · be48ce7 · be48ce7
1 parent 94bce4a
commit be48ce7
Showing 1 changed file with 10 additions and 1 deletion.
diff --git a/extensions/tokenizers/rust/src/models/bert.rs b/extensions/tokenizers/rust/src/models/bert.rs
@@ -492,7 +492,16 @@ impl BertClassificationHead {
             Ok(layer) => Some(layer),
             Err(_) => None,
         };
-        let output = Linear::load(vb.pp("classifier"), config.hidden_size, n_classes, None)?;
+        let output = match Linear::load(vb.pp("classifier"), config.hidden_size, n_classes, None) {
+            Ok(output) => output,
+            Err(err) => {
+                if let Ok(output) = Linear::load(vb, config.hidden_size, n_classes, None) {
+                    output
+                } else {
+                    return Err(err);
+                }
+            }
+        };
 
         Ok(Self {
             pooler,