avoid unnecessary unsqueeze

open-spaced-repetition · Oct 29, 2024 · 57c562b · 57c562b
1 parent 151fddc
commit 57c562b
Showing 1 changed file with 6 additions and 8 deletions.
diff --git a/src/dataset.rs b/src/dataset.rs
@@ -106,26 +106,24 @@ impl<B: Backend> Batcher<FSRSItem, FSRSBatch<B>> for FSRSBatcher<B> {
                     item.history().map(|r| (r.delta_t, r.rating)).unzip();
                 delta_t.resize(pad_size, 0);
                 rating.resize(pad_size, 0);
-                let delta_t = Tensor::<B, 1>::from_floats(
+                let delta_t = Tensor::<B, 2>::from_floats(
                     TensorData::new(
                         delta_t,
                         Shape {
-                            dims: vec![pad_size],
+                            dims: vec![1, pad_size],
                         },
                     ),
                     &self.device,
-                )
-                .unsqueeze();
-                let rating = Tensor::<B, 1>::from_data(
+                );
+                let rating = Tensor::<B, 2>::from_data(
                     TensorData::new(
                         rating,
                         Shape {
-                            dims: vec![pad_size],
+                            dims: vec![1, pad_size],
                         },
                     ),
                     &self.device,
-                )
-                .unsqueeze();
+                );
                 (delta_t, rating)
             })
             .unzip();