fix metric output reshaping

don't use batch dim

Author: radka-j

autoemulate/core/metrics.py

@@ -373,7 +373,7 @@ def __call__(
                         "distributions."
                     )
                     raise ValueError(msg)
-                return model_nll_output.reshape(*y_true.shape)
+                return model_nll_output.reshape(*y_true.shape[1:])
             msg = (
                 f"Unknown reduction method: {metric_params.reduction}. "
                 "Expected 'mean' or 'none'."
@@ -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).reshape(*y_true.shape)
+            return (model_nll_output - trivial_nll_output).reshape(*y_true.shape[1:])
         msg = (
             f"Unknown reduction method: {metric_params.reduction}. "
             "Expected 'mean' or 'none'."