Remove the TorchANI import from the top level (openmm#44)

raimis · Dec 10, 2021 · 93a7836 · 93a7836
1 parent a830fa1
commit 93a7836
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/src/pytorch/EnergyShifter.py b/src/pytorch/EnergyShifter.py
@@ -23,13 +23,17 @@
 
 import torch
 from torch import Tensor
-from torchani.nn import SpeciesConverter
-from torchani.utils import EnergyShifter, SpeciesEnergies
-from typing import Optional, Tuple
+from typing import NamedTuple, Optional, Tuple
 
+class SpeciesEnergies(NamedTuple):
+    species: Tensor
+    energies: Tensor
 
 class TorchANIEnergyShifter(torch.nn.Module):
 
+    from torchani.nn import SpeciesConverter # https://github.com/openmm/NNPOps/issues/44
+    from torchani.utils import EnergyShifter # https://github.com/openmm/NNPOps/issues/44
+
     def __init__(self, converter: SpeciesConverter, shifter: EnergyShifter, atomicNumbers: Tensor) -> None:
 
         super().__init__()