Commit 9715749f authored by Nan Xu's avatar Nan Xu Committed by Farhad Ramezanghorbani
Browse files

Fix the dimension in self_energies for dataset containing only one elements (#302)

* Update utils.py

if I train molecule like Oxygen which contains only one element, the self.self_energies is a scalar and will make the operation self_energies = self.self_energies[species] crash. Besides, the species of Oxygen is [[3,3]] and is over the boundary of the energy list, so need to minus the minimum.

* Update utils.py
parent 379d2b33
......@@ -172,7 +172,7 @@ class EnergyShifter(torch.nn.Module):
X = torch.cat((X, torch.ones(X.shape[0], 1).to(torch.double)), dim=-1)
y = energies.unsqueeze(dim=-1)
coeff_, _, _, _ = np.linalg.lstsq(X, y, rcond=None)
return coeff_.squeeze()
return coeff_.squeeze(-1)
def sae(self, species):
"""Compute self energies for molecules.
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment