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