Remove TorchANI imports for type annotations

.github/workflows/ci.yml

@@ -33,14 +33,14 @@ jobs:
             pytorch: "2.0.*"
 
           # Latest supported versions (with CUDA)
-          - name: Linux (CUDA 12, Python 3.13, PyTorch 2.5)
+          - name: Linux (CUDA 12, Python 3.13, PyTorch >=2.0)
             enable_cuda: true
             cuda: "12.6.0"
             gcc: "10.3.*"
             nvcc: "12.*"
             python: "3.13.*"
             torchani: "2.2.*"
-            pytorch: "2.5.*"
+            pytorch: "2.*"
 
     steps:
     - name: Check out

src/pytorch/BatchedNN.py

@@ -36,9 +36,7 @@ class SpeciesEnergies(NamedTuple):
 
 class _BatchedNN(torch.nn.Module):
 
-    from torchani.nn import ANIModel, Ensemble, SpeciesConverter # https://github.com/openmm/NNPOps/issues/44
-
-    def __init__(self, converter: SpeciesConverter, ensemble: Union[ANIModel, Ensemble], atomicNumbers: Tensor):
+    def __init__(self, converter, ensemble, atomicNumbers: Tensor):
 
         super().__init__()
 
@@ -113,9 +111,7 @@ def forward(self, species_aev: Tuple[Tensor, Tensor]) -> SpeciesEnergies:
 
 class TorchANIBatchedNN(torch.nn.ModuleList):
 
-    from torchani.nn import ANIModel, Ensemble, SpeciesConverter # https://github.com/openmm/NNPOps/issues/44
-
-    def __init__(self, converter: SpeciesConverter, ensemble: Union[ANIModel, Ensemble], atomicNumbers: Tensor):
+    def __init__(self, converter, ensemble, atomicNumbers: Tensor):
         super().__init__([_BatchedNN(converter, ensemble, atomicNumbers)])
 
     def forward(self, species_aev: Tuple[Tensor, Tensor]) -> SpeciesEnergies:

src/pytorch/EnergyShifter.py

@@ -31,10 +31,7 @@ class SpeciesEnergies(NamedTuple):
 
 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:
+    def __init__(self, converter, shifter, atomicNumbers: Tensor) -> None:
 
         super().__init__()
 
@@ -50,4 +47,4 @@ def forward(self, species_energies: Tuple[Tensor, Tensor],
 
         species, energies = species_energies
 
-        return SpeciesEnergies(species, energies + self.self_energies)
+        return SpeciesEnergies(species, energies + self.self_energies)

src/pytorch/SpeciesConverter.py

@@ -23,9 +23,7 @@ class SpeciesCoordinates(NamedTuple):
 
 class TorchANISpeciesConverter(torch.nn.Module):
 
-    from torchani.nn import SpeciesConverter
-
-    def __init__(self, converter: SpeciesConverter, atomicNumbers: Tensor) -> None:
+    def __init__(self, converter, atomicNumbers: Tensor) -> None:
 
         super().__init__()
 

src/pytorch/SymmetryFunctions.py

@@ -60,10 +60,7 @@ class TorchANISymmetryFunctions(torch.nn.Module):
         >>> print(energy, forces)
     """
 
-    from torchani import AEVComputer # https://github.com/openmm/NNPOps/pull/38
-    from torchani import SpeciesConverter # https://github.com/openmm/NNPOps/pull/38
-
-    def __init__(self, converter: SpeciesConverter, symmFunc: AEVComputer, atomicNumbers: Tensor) -> None:
+    def __init__(self, converter, symmFunc, atomicNumbers: Tensor) -> None:
         """
         Arguments:
             converter: an instance of torchani.nn.SpeciesConverter (https://aiqm.github.io/torchani/api.html#torchani.SpeciesConverter)