Return full state as embedding if embed layer requested is larger tha…

…n model's layer count.
Ai00-X · Jan 10, 2025 · 11954e0 · 11954e0
1 parent 5ae9de6
commit 11954e0
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/crates/ai00-core/src/run.rs b/crates/ai00-core/src/run.rs
@@ -712,9 +712,11 @@ impl Runtime {
 
             let backed = self.state.back(batch).await?;
             if let GenerateKind::Embed { layer } = context.request.kind {
-                let layer = layer.clamp(0, self.info.num_layer - 1);
                 let backed = backed.clone();
-                let embed = self.state.embed(layer, backed)?.to_vec();
+                let embed = match layer {
+                    x if x < self.info.num_layer => self.state.embed(layer, backed)?.to_vec(),
+                    _ => backed.to_vec(),
+                };
                 let _ = context.sender.send(Token::Embed(embed));
             }