Update utils.py (#338)

aiqm · Oct 15, 2019 · 18e63f1 · 18e63f1
1 parent 053f086
commit 18e63f1
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/torchani/utils.py b/torchani/utils.py
@@ -215,8 +215,8 @@ def forward(self, species_energies):
         """(species, molecular energies)->(species, molecular energies + sae)
         """
         species, energies = species_energies
-        sae = self.sae(species).to(energies.dtype).to(energies.device)
-        return species, energies + sae
+        sae = self.sae(species).to(energies.device)
+        return species, energies.to(sae.dtype) + sae
 
 
 class ChemicalSymbolsToInts: