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

raimis · Dec 10, 2021 · 2634298 · 2634298
1 parent 074e48e
commit 2634298
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/src/pytorch/SpeciesConverter.py b/src/pytorch/SpeciesConverter.py
@@ -15,11 +15,12 @@
 
 import torch
 from torch import Tensor
-from torchani.nn import SpeciesCoordinates, SpeciesConverter
 from typing import Optional, Tuple
 
 class TorchANISpeciesConverter(torch.nn.Module):
 
+    from torchani.nn import SpeciesConverter
+
     def __init__(self, converter: SpeciesConverter, atomicNumbers: Tensor) -> None:
 
         super().__init__()
@@ -32,8 +33,8 @@ def __init__(self, converter: SpeciesConverter, atomicNumbers: Tensor) -> None:
 
     def forward(self, species_coordinates: Tuple[Tensor, Tensor],
                 cell: Optional[Tensor] = None,
-                pbc: Optional[Tensor] = None) -> SpeciesCoordinates:
+                pbc: Optional[Tensor] = None) -> Tuple[Tensor, Tensor]:
 
         _, coordinates = species_coordinates
 
-        return SpeciesCoordinates(self.species, coordinates)
+        return self.species, coordinates