correct shape of logits

open-spaced-repetition · Aug 25, 2023 · 06a9863 · 06a9863
1 parent e25a7a8
commit 06a9863
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/training.rs b/src/training.rs
@@ -36,7 +36,7 @@ impl<B: Backend<FloatElem = f32>> Model<B> {
  stability.clone(),
  );
  let logits =
- Tensor::cat(vec![-retention.clone() + 1, retention.clone()], 0).reshape([2, -1]);
+ Tensor::cat(vec![-retention.clone() + 1, retention.clone()], 1);
  info!("stability: {}", &stability);
  info!(
  "delta_ts: {}",