change metric wrappers to keep graddable metrics graddable

mackelab · Feb 7, 2024 · fdaf384 · fdaf384
1 parent 551b9b7
commit fdaf384
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/labproject/metrics/utils.py b/labproject/metrics/utils.py
@@ -24,7 +24,8 @@ def wrapper(*args, **kwargs):
             metric = func(*args, **kwargs)
 
             # Convert output to tensor
-            metric = torch.tensor(metric)
+            if not isinstance(metric, torch.Tensor):
+                metric = torch.tensor(metric)
 
             return metric