Update autoemulate/core/metrics.py

Co-authored-by: Sam Greenbury <[email protected]>

Author: radka-j

autoemulate/core/metrics.py

@@ -403,7 +403,7 @@ def __call__(
             if model_nll_output is None:
                 msg = "Per-output MLL not available for non-Independent distributions."
                 raise ValueError(msg)
-            return model_nll_output - trivial_nll_output
+            return (model_nll_output - trivial_nll_output).reshape(*y_true.shape)
         msg = (
             f"Unknown reduction method: {metric_params.reduction}. "
             "Expected 'mean' or 'none'."