Fix loss accumulation

src/base_tester.py

@@ -68,8 +68,8 @@ def _test_iteration(
         Returns:
             torch.Tensor: The loss for the batch.
         """
-        # x, y = batch  # type: ignore
-        # y_hat = self._model(x)
+        x, y = batch  # type: ignore # noqa
+        y_hat = self._model(x)  # type: ignore # noqa
         # TODO: Compute your metrics here!
         return {}
 

src/base_trainer.py

@@ -102,7 +102,7 @@ def _train_val_iteration(
         x, y = batch
         y_hat = self._model(x)
         losses: Dict[str, torch.Tensor] = self._training_loss(y, y_hat)
-        loss = sum([v for v in losses.values()])
+        loss: torch.Tensor = torch.sum(torch.tensor([v for v in losses.values()]))
         return loss, losses
 
     def _train_epoch(

src/losses/mse.py

@@ -10,14 +10,18 @@
 """
 
 
+from typing import Dict
+
 import torch
 
 
 class MSELoss:
     def __init__(self, reduction: str):
         self._reduction = reduction
 
-    def __call__(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
+    def __call__(
+        self, y_pred: torch.Tensor, y_true: torch.Tensor
+    ) -> Dict[str, torch.Tensor]:
         return {
             "mse": torch.nn.functional.mse_loss(
                 y_pred, y_true, reduction=self._reduction