indexing by column

Author: FNTwin

src/openqdc/datasets/base.py

@@ -289,7 +289,7 @@ def __getitem__(self, idx: int):
             positions=positions,
             atomic_numbers=z,
             charges=c,
-            e0=self.__isolated_atom_energies__[..., z, c + shift],
+            e0=self.__isolated_atom_energies__[..., z, c + shift].T,
             energies=energies,
             name=name,
             subset=subset,