Fix metrics

santiag0m · May 16, 2022 · 9611ea7 · 9611ea7
1 parent 2265bf2
commit 9611ea7
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/lib/utils/metrics.py b/lib/utils/metrics.py
@@ -19,7 +19,7 @@ def compute_accuracy(model: nn.Module, dataloader: DataLoader) -> torch.Tensor:
             preds = torch.argmax(preds, dim=-1)
             correct = preds == targets
             accuracy.update(correct.cpu())
-    return accuracy.value
+    return accuracy.value.item()
 
 
 def compute_mse(
@@ -39,7 +39,7 @@ def compute_mse(
             batch_mse = (targets - preds) ** 2
             batch_mse = batch_mse.sum(axis=1)
             mse.update(batch_mse.cpu())
-    return mse.value
+    return mse.value.item()
 
 
 def compute_bias(