fix predictions to not use argmax

models/src/models/focus_module.py

@@ -86,7 +86,7 @@ class FocusLitModule(LightningModule):
         y = batch["focus_value"]
         logits = self.forward(x)
         loss = self.criterion(logits, y)
-        preds = torch.argmax(logits, dim=1)
+        preds = torch.squeeze(logits)
         return loss, preds, y
 
     def training_step(self, batch: Any, batch_idx: int):