Commit d32081e9 authored by Gao, Xiang's avatar Gao, Xiang Committed by Farhad Ramezanghorbani
Browse files

Use modern type annotations for utils.py (#375)

* Use modern type annotations for utils.py

* update (#377)

* Remove python2 support from CI, _six, etc. (#370)

* Remove python2 step 1

* badge

* Delete _six.py

* Simplify triple_by_molecule (#368)

* Simplify triple_by_molecule

* fix

* fix

* fix

* Use @ syntax (#371)

* Use modern type annotations for nn.py (#373)

* Use modern type annotations for nn.py

* flake8

* Use modern type annotations for models.py (#374)
parent 98bb1237
import torch
from torch import Tensor
import torch.utils.data
import math
import numpy as np
......@@ -210,8 +211,7 @@ class EnergyShifter(torch.nn.Module):
properties['energies'] = energies
return atomic_properties, properties
def forward(self, species_energies):
# type: (Tuple[torch.Tensor, torch.Tensor]) -> Tuple[torch.Tensor, torch.Tensor]
def forward(self, species_energies: Tuple[Tensor, Tensor]) -> Tuple[Tensor, Tensor]:
"""(species, molecular energies)->(species, molecular energies + sae)
"""
species, energies = species_energies
......
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