fix: correcting model output validation for scalar energies

IntelLabs · Nov 27, 2024 · a9c3c41 · a9c3c41
1 parent 7db4282
commit a9c3c41
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/matsciml/common/types.py b/matsciml/common/types.py
@@ -141,8 +141,12 @@ def standardize_total_energy(
             ``ValueError``.
         """
         if isinstance(values, torch.Tensor):
-            # drop all redundant dimensions
-            values = values.squeeze()
+            if values.ndim > 1:
+                # drop all redundant dimensions
+                values = values.squeeze()
+            # we want to at least have a vector of energies
+            if values.ndim == 0:
+                values = values.unsqueeze(0)
             # last step is an assertion check for QA
             if values.ndim != 1:
                 raise ValueError(