fixed the parse_reduction to take values from min, max, median

graphium/nn/architectures/global_architectures.py

@@ -599,16 +599,23 @@ def _parse_reduction(self, reduction: Optional[Union[str, Callable]]) -> Optiona
         elif reduction == "sum":
             return torch.sum
         elif reduction == "max":
-            return torch.max
+            def max_vals(x, dim):
+                return torch.max(x, dim=dim).values
+            return max_vals
         elif reduction == "min":
-            return torch.min
+            def min_vals(x, dim):
+                return torch.min(x, dim=dim).values
+            return min_vals
         elif reduction == "median":
-            return torch.median
+            def median_vals(x, dim):
+                return torch.median(x, dim=dim).values
+            return median_vals
         elif callable(reduction):
             return reduction
         else:
             raise ValueError(f"Unknown reduction {reduction}")
 
+
     def _parse_layers(self, layer_type, residual_type):
         # Parse the layer and residuals
         from graphium.utils.spaces import ENSEMBLE_LAYERS_DICT, RESIDUALS_DICT

graphium/nn/ensemble_layers.py

@@ -391,11 +391,17 @@ def _parse_reduction(self, reduction: Optional[Union[str, Callable]]) -> Optiona
         elif reduction == "sum":
             return torch.sum
         elif reduction == "max":
-            return torch.max
+            def max_vals(x, dim):
+                return torch.max(x, dim=dim).values
+            return max_vals
         elif reduction == "min":
-            return torch.min
+            def min_vals(x, dim):
+                return torch.min(x, dim=dim).values
+            return min_vals
         elif reduction == "median":
-            return torch.median
+            def median_vals(x, dim):
+                return torch.median(x, dim=dim).values
+            return median_vals
         elif callable(reduction):
             return reduction
         else: