refactor: updating shape checking for MaceEnergyForceTask as well

Author: Lee, Kin Long Kelvin

matsciml/models/base.py

@@ -1584,7 +1584,16 @@ def _compute_losses(
             else:
                 coefficient = self.loss_coeff[key]
 
-            losses[key] = self.loss_func(predictions[key], target_val) * (
+            model_outputs = predictions[key]
+            # attempt to reshape model outputs for correct broadcasting
+            if model_outputs.shape != target_val.shape:
+                try:
+                    model_outputs = model_outputs.reshape(target_val.shape)
+                except RuntimeError as e:
+                    raise RuntimeError(
+                        f"Unable to reconile shapes for preds/labels; preds: {model_outputs.shape}, labels: {target_val.shape}."
+                    ) from e
+            losses[key] = self.loss_func(model_outputs, target_val) * (
                 coefficient / predictions[key].numel()
             )
         total_loss: torch.Tensor = sum(losses.values())