Merge pull request #63 from SimonBoothroyd/dipole_shape

Dipole reshape
SimonBoothroyd · May 30, 2023 · 6c6cf2e · 6c6cf2e
2 parents ba61333 + 683980f
commit 6c6cf2e
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/nagl/training/loss.py b/nagl/training/loss.py
@@ -164,7 +164,8 @@ def evaluate_loss(
         prediction: typing.Dict[str, torch.Tensor],
     ) -> torch.Tensor:
         metric_func = get_metric(self.metric)
-        target_dipole = labels[self.dipole_column].squeeze()
+        # reshape as it can be flat
+        target_dipole = torch.reshape(labels[self.dipole_column], (-1, 3))
         n_atoms_per_molecule = (
             (molecules.n_atoms,)
             if isinstance(molecules, DGLMolecule)