fix: ClassificationHead

kozistr · Nov 17, 2024 · fd6dcce · fd6dcce
1 parent 751dc02
commit fd6dcce
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/backends/candle/src/models/flash_gte.rs b/backends/candle/src/models/flash_gte.rs
@@ -1,7 +1,8 @@
 use crate::flash_attn::flash_attn_varlen;
 use crate::layers::{HiddenAct, LayerNorm, Linear};
-use crate::models::{
-    GTEClassificationHead, GTEConfig, Model, NTKScaling, PositionEmbeddingType, RopeScaling,
+use crate::models::gte::{
+    ClassificationHead, GTEClassificationHead, GTEConfig, Model, NTKScaling, PositionEmbeddingType,
+    RopeScaling,
 };
 use candle::{DType, Device, IndexOp, Result, Tensor};
 use candle_nn::{Embedding, Module, VarBuilder};