edit dim

src/sparcscore/ml/plmodels.py

@@ -205,7 +205,7 @@ def configure_optimizers(self):
     def training_step(self, batch):
         data, target = batch
         print("target training ", target)
-        target = target.unsqueeze(0) # Add dimension for regression
+        target = target.unsqueeze(1) # Add dimension for regression
         output = self.network(data) # Forward pass, only one output
         loss = F.mse_loss(output, target) # L2 loss
 
@@ -219,7 +219,7 @@ def training_step(self, batch):
 
     def validation_step(self, batch):
         data, target = batch
-        target = target.unsqueeze(0)
+        target = target.unsqueeze(1)
         print("target val ", target)
 
         output = self.network(data)