Updated sentence classification to work with single input
@@ -317,8 +317,8 @@ impl SequenceClassificationModel {
false);
output.softmax(-1, Kind::Float).detach().to(Device::Cpu)
});
- let label_indices = output.as_ref().argmax(-1, true).squeeze();
- let scores = output.gather(1, &label_indices.unsqueeze(1), false).squeeze();
+ let label_indices = output.as_ref().argmax(-1, true).squeeze1(1);
+ let scores = output.gather(1, &label_indices.unsqueeze(-1), false).squeeze1(1);
let label_indices = label_indices.iter::<i64>().unwrap().collect::<Vec<i64>>();
let scores = scores.iter::<f64>().unwrap().collect::<Vec<f64>>();
