Fix simple regression batch targets (#2379)

tracel-ai · Oct 17, 2024 · 296c526 · 296c526
1 parent 359ad08
commit 296c526
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/examples/simple-regression/src/model.rs b/examples/simple-regression/src/model.rs
@@ -50,7 +50,7 @@ impl<B: Backend> RegressionModel<B> {
     }
 
     pub fn forward_step(&self, item: DiabetesBatch<B>) -> RegressionOutput<B> {
-        let targets: Tensor<B, 2> = item.targets.unsqueeze();
+        let targets: Tensor<B, 2> = item.targets.unsqueeze_dim(1);
         let output: Tensor<B, 2> = self.forward(item.inputs);
 
         let loss = MseLoss::new().forward(output.clone(), targets.clone(), Mean);