Skip to content

GitLab

You need to sign in or sign up before continuing.
Commit 7499c8d4 authored by Gao, Xiang's avatar Gao, Xiang Committed by Farhad Ramezanghorbani
Browse files

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

parent 2c0e6df5
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 2 deletions
torchani/models.py
View file @ 7499c8d4
...@@ -28,6 +28,7 @@ shouldn't be used anymore. ...@@ -28,6 +28,7 @@ shouldn't be used anymore.
""" """
import torch import torch
from torch import Tensor
from typing import Tuple from typing import Tuple
from pkg_resources import resource_filename from pkg_resources import resource_filename
from . import neurochem from . import neurochem
...@@ -88,8 +89,7 @@ class BuiltinNet(torch.nn.Module): ...@@ -88,8 +89,7 @@ class BuiltinNet(torch.nn.Module):
self.neural_networks = neurochem.load_model_ensemble( self.neural_networks = neurochem.load_model_ensemble(
self.species, self.ensemble_prefix, self.ensemble_size) self.species, self.ensemble_prefix, self.ensemble_size)
def forward(self, species_coordinates): def forward(self, species_coordinates: Tuple[Tensor, Tensor]) -> Tuple[Tensor, Tensor]:
# type: (Tuple[torch.Tensor, torch.Tensor]) -> Tuple[torch.Tensor, torch.Tensor]
"""Calculates predicted properties for minibatch of configurations """Calculates predicted properties for minibatch of configurations
Args: Args:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment