energies.ndim=1

Author: Ewan Wallace

src/neural_optimiser/calculators/_mmff94.py

@@ -93,7 +93,7 @@ def _calculate(self, batch: Data | Batch) -> tuple[torch.Tensor, torch.Tensor]:
         mol = self._prepare_mol(mol)
 
         ff = self._build_forcefield(mol)
-        energy = torch.tensor(ff.CalcEnergy())
+        energy = torch.tensor([ff.CalcEnergy()])
         grad = (
             torch.Tensor(ff.CalcGrad()) * -KCAL_PER_MOL_TO_EV
         )  # Convert to eV/Å for convergence criteria
@@ -106,4 +106,4 @@ def get_energies(self, batch: Data | Batch) -> torch.Tensor:
         mol = self._prepare_mol(mol)
 
         ff = self._build_forcefield(mol)
-        return torch.tensor(ff.CalcEnergy())
+        return torch.tensor([ff.CalcEnergy()])

tests/calculators/test_mmff94.py

@@ -48,7 +48,7 @@ def test_single_conformer_energy_and_forces_shape(batch):
 
     # energy is a scalar tensor
     assert isinstance(energy, torch.Tensor)
-    assert energy.ndim == 0
+    assert energy.ndim == 1
 
     # forces match [n_atoms, 3]
     assert isinstance(forces, torch.Tensor)