Add code so loss function uses torch.Size([x,1]) instead of torch.Size([x])

This was done to prevent the error message:
"Using a target size (torch.Size([64])) that is different to the input size (torch.Size([64, 1]))".

src/models/focus_module.py

@@ -85,7 +85,7 @@ class FocusLitModule(LightningModule):
         x = batch["image"]
         y = batch["focus_value"]
         logits = self.forward(x)
-        loss = self.criterion(logits, y)
+        loss = self.criterion(logits, y.unsqueeze(1))
         preds = torch.squeeze(logits)
         return loss, preds, y
 
@@ -210,7 +210,7 @@ class FocusMSELitModule(LightningModule):
         x = batch["image"]
         y = batch["focus_value"]
         logits = self.forward(x)
-        loss = self.criterion(logits, y)
+        loss = self.criterion(logits, y.unsqueeze(1))
         preds = torch.squeeze(logits)
         return loss, preds, y