Move all torch.ops.load calls to the __init__.py scripts

CMakeLists.txt

@@ -77,8 +77,6 @@ install(FILES src/pytorch/__init__.py
               src/pytorch/OptimizedTorchANI.py
               src/pytorch/SpeciesConverter.py
               src/pytorch/SymmetryFunctions.py
-              src/pytorch/neighbors/__init__.py
-              src/pytorch/neighbors/getNeighborPairs.py
         DESTINATION ${Python3_SITEARCH}/${NAME})
 install(FILES src/pytorch/neighbors/__init__.py
               src/pytorch/neighbors/getNeighborPairs.py

src/pytorch/BatchedNN.py

@@ -21,14 +21,12 @@
 # SOFTWARE.
 #
 
-import os
 import torch
 from torch import nn
 from torch import Tensor
 from torch.nn import functional as F
 from typing import List, NamedTuple, Tuple, Union
 
-torch.ops.load_library(os.path.join(os.path.dirname(__file__), 'libNNPOpsPyTorch.so'))
 batchedLinear = torch.ops.NNPOpsBatchedNN.BatchedLinear
 
 

src/pytorch/CFConv.py

@@ -21,15 +21,11 @@
 # SOFTWARE.
 #
 
-import os.path
 import torch
 from torch import Tensor
 
 from NNPOps.CFConvNeighbors import CFConvNeighbors
 
-torch.ops.load_library(os.path.join(os.path.dirname(__file__), 'libNNPOpsPyTorch.so'))
-torch.classes.load_library(os.path.join(os.path.dirname(__file__), 'libNNPOpsPyTorch.so'))
-
 class CFConv(torch.nn.Module):
     """
     Optimized continious-filter convolution layer (CFConv)
@@ -84,4 +80,4 @@ def __init__(self, gaussianWidth: float, activation: str,
 
     def forward(self, neighbors: CFConvNeighbors, positions: Tensor, input: Tensor) -> Tensor:
 
-        return CFConv.operation(self.holder, neighbors.holder, positions, input)
+        return CFConv.operation(self.holder, neighbors.holder, positions, input)

src/pytorch/CFConvNeighbors.py

@@ -21,12 +21,9 @@
 # SOFTWARE.
 #
 
-import os.path
 import torch
 from torch import Tensor
 
-torch.classes.load_library(os.path.join(os.path.dirname(__file__), 'libNNPOpsPyTorch.so'))
-
 class CFConvNeighbors(torch.nn.Module):
     """
     Optimized nearest-neighbor implementation for the continious-filter convolution (CFConf)
@@ -45,4 +42,4 @@ def __init__(self, cutoff: float) -> None:
     @torch.jit.export
     def build(self, positions: Tensor) -> None:
 
-        self.holder.build(positions)
+        self.holder.build(positions)

src/pytorch/SymmetryFunctions.py

@@ -21,14 +21,10 @@
 # SOFTWARE.
 #
 
-import os.path
 from typing import List, Optional, Tuple
 import torch
 from torch import Tensor
 
-torch.ops.load_library(os.path.join(os.path.dirname(__file__), 'libNNPOpsPyTorch.so'))
-torch.classes.load_library(os.path.join(os.path.dirname(__file__), 'libNNPOpsPyTorch.so'))
-
 Holder = torch.classes.NNPOpsANISymmetryFunctions.Holder
 operation = torch.ops.NNPOpsANISymmetryFunctions.operation
 
@@ -124,4 +120,4 @@ def forward(self, species_positions: Tuple[Tensor, Tensor],
         radial, angular = operation(self.holder, positions[0], cell)
         features = torch.cat((radial, angular), dim=1).unsqueeze(0)
 
-        return species, features
+        return species, features

src/pytorch/__init__.py

@@ -1,5 +1,10 @@
 '''
 High-performance PyTorch operations for neural network potentials
 '''
+import os.path
+import torch
 
-from NNPOps.OptimizedTorchANI import OptimizedTorchANI
+torch.ops.load_library(os.path.join(os.path.dirname(__file__), 'libNNPOpsPyTorch.so'))
+torch.classes.load_library(os.path.join(os.path.dirname(__file__), 'libNNPOpsPyTorch.so'))
+
+from NNPOps.OptimizedTorchANI import OptimizedTorchANI

src/pytorch/neighbors/__init__.py

@@ -1,5 +1,9 @@
 '''
 Neighbor operations
 '''
+import os.path
+import torch
 
-from NNPOps.neighbors.getNeighborPairs import getNeighborPairs
+torch.ops.load_library(os.path.join(os.path.dirname(os.path.dirname(__file__)), 'libNNPOpsPyTorch.so'))
+
+from NNPOps.neighbors.getNeighborPairs import getNeighborPairs