Use modern type annotations for models.py (#374)

aiqm · Nov 7, 2019 · 7499c8d · 7499c8d
1 parent 2c0e6df
commit 7499c8d
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/torchani/models.py b/torchani/models.py
@@ -28,6 +28,7 @@
 """
 
 import torch
+from torch import Tensor
 from typing import Tuple
 from pkg_resources import resource_filename
 from . import neurochem
@@ -88,8 +89,7 @@ def __init__(self, info_file):
         self.neural_networks = neurochem.load_model_ensemble(
             self.species, self.ensemble_prefix, self.ensemble_size)
 
-    def forward(self, species_coordinates):
-        # type: (Tuple[torch.Tensor, torch.Tensor]) -> Tuple[torch.Tensor, torch.Tensor]
+    def forward(self, species_coordinates: Tuple[Tensor, Tensor]) -> Tuple[Tensor, Tensor]:
         """Calculates predicted properties for minibatch of configurations
 
         Args: