[Fix #648] Pass cell and pbc as args when calling AEVComputer so forward hook works via TorchScript

tests/test_aev_hook.py

@@ -0,0 +1,49 @@
+import unittest
+
+import torch
+import torchani
+
+from torch import Tensor
+from typing import Optional, Tuple
+
+from torchani.testing import TestCase
+
+
+class TestAEVHook(TestCase):
+    def test_aev_hook(self):
+        # Create a test module.
+
+        class TestModule(torch.nn.Module):
+            def __init__(self):
+                super(TestModule, self).__init__()
+
+                # Create a ANI2x model.
+                self._model = torchani.models.ANI2x()
+
+                # Define a dummy hook. This doesn't do anything with the output
+                # but triggers the exception when kwargs are used when calling
+                # AEVComputer.forward.
+                def hook(
+                    module,
+                    input: Tuple[
+                        Tuple[Tensor, Tensor], Optional[Tensor], Optional[Tensor]
+                    ],
+                    output: torchani.aev.SpeciesAEV,
+                ):
+                    pass
+
+                # Register the hook.
+                self._aev_hook = self._model.aev_computer.register_forward_hook(hook)
+
+            def forward(self, species, coordinates):
+                return self._model((species, coordinates))
+
+        # Create a test module.
+        model = TestModule()
+
+        # Convert the model to TorchScript.
+        model = torch.jit.script(model)
+
+
+if __name__ == "__main__":
+    unittest.main()

torchani/models.py

@@ -103,7 +103,7 @@ def forward(self, species_coordinates: Tuple[Tensor, Tensor],
         if species_coordinates[0].ge(self.aev_computer.num_species).any():
             raise ValueError(f'Unknown species found in {species_coordinates[0]}')
 
-        species_aevs = self.aev_computer(species_coordinates, cell=cell, pbc=pbc)
+        species_aevs = self.aev_computer(species_coordinates, cell, pbc)
         species_energies = self.neural_networks(species_aevs)
         return self.energy_shifter(species_energies)
 
@@ -135,7 +135,7 @@ def atomic_energies(self, species_coordinates: Tuple[Tensor, Tensor],
         """
         if self.periodic_table_index:
             species_coordinates = self.species_converter(species_coordinates)
-        species, aevs = self.aev_computer(species_coordinates, cell=cell, pbc=pbc)
+        species, aevs = self.aev_computer(species_coordinates, cell, pbc)
         atomic_energies = self.neural_networks._atomic_energies((species, aevs))
         self_energies = self.energy_shifter.self_energies.clone().to(species.device)
         self_energies = self_energies[species]
@@ -236,7 +236,7 @@ def atomic_energies(self, species_coordinates: Tuple[Tensor, Tensor],
         """
         if self.periodic_table_index:
             species_coordinates = self.species_converter(species_coordinates)
-        species, aevs = self.aev_computer(species_coordinates, cell=cell, pbc=pbc)
+        species, aevs = self.aev_computer(species_coordinates, cell, pbc)
         members_list = []
         for nnp in self.neural_networks:
             members_list.append(nnp._atomic_energies((species, aevs)).unsqueeze(0))
@@ -322,7 +322,7 @@ def members_energies(self, species_coordinates: Tuple[Tensor, Tensor],
         """
         if self.periodic_table_index:
             species_coordinates = self.species_converter(species_coordinates)
-        species, aevs = self.aev_computer(species_coordinates, cell=cell, pbc=pbc)
+        species, aevs = self.aev_computer(species_coordinates, cell, pbc)
         member_outputs = []
         for nnp in self.neural_networks:
             unshifted_energies = nnp((species, aevs)).energies