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
...@@ -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