[rust] Fix bert model classifier loading

deepjavalibrary · Aug 26, 2024 · d5929c0 · d5929c0
1 parent 94bce4a
commit d5929c0
Showing 1 changed file with 11 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,17 @@ 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,