From 45a671cf52cab48f16384c89ebaaaa52558d87bc Mon Sep 17 00:00:00 2001
From: tuna2134 <hello@tuna2134.dev>
Date: Fri, 18 Oct 2024 13:35:23 +0000
Subject: [PATCH] fix compile

---
 sbv2_core/src/bert.rs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/sbv2_core/src/bert.rs b/sbv2_core/src/bert.rs
index cf0bdcf..0710bf6 100644
--- a/sbv2_core/src/bert.rs
+++ b/sbv2_core/src/bert.rs
@@ -1,5 +1,5 @@
 use crate::error::Result;
-use ndarray::{Array2, Ix3};
+use ndarray::{Array2, Ix2};
 use ort::Session;
 
 pub fn predict(
@@ -16,7 +16,7 @@ pub fn predict(
 
     let output = outputs["output"]
         .try_extract_tensor::<f32>()?
-        .into_dimensionality::<Ix3>()?
+        .into_dimensionality::<Ix2>()?
         .to_owned();
 
     Ok(output)