diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 50d797e9..00a43807 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -25,23 +25,32 @@ jobs:
           # Oldest supported versions
           - name: Linux (CUDA 11.8, Python 3.10, PyTorch 2.0)
             enable_cuda: true
+            set_torch_dir: true
             cuda: "11.8.0"
             gcc: "10.3.*"
             nvcc: "11.8"
             python: "3.10.*"
-            torchani: "2.2.*"
             pytorch: "2.0.*"
 
-          # Latest supported versions (with CUDA)
-          - name: Linux (CUDA 12, Python 3.13, PyTorch 2.5)
+          - name: Linux (CUDA 12.6, Python 3.13, PyTorch 2.5)
             enable_cuda: true
+            set_torch_dir: true
             cuda: "12.6.0"
             gcc: "10.3.*"
             nvcc: "12.*"
             python: "3.13.*"
-            torchani: "2.2.*"
             pytorch: "2.5.*"
 
+          # Newest supported versions
+          - name: Linux (CUDA 12.9, Python 3.13, PyTorch 2.9)
+            enable_cuda: true
+            set_torch_dir: false
+            cuda: "12.9.0"
+            gcc: "14.*"
+            nvcc: "12.9"
+            python: "3.13.*"
+            pytorch: "2.9.*"
+
     steps:
     - name: Check out
       uses: actions/checkout@v2
@@ -50,7 +59,7 @@ jobs:
       uses: jlumbroso/free-disk-space@main
 
     - name: Install CUDA Toolkit
-      uses: Jimver/cuda-toolkit@v0.2.21
+      uses: Jimver/cuda-toolkit@v0.2.30
       with:
         cuda: ${{ matrix.cuda }}
         linux-local-args: '["--toolkit", "--override"]'
@@ -68,7 +77,6 @@ jobs:
     - name: Prepare dependencies
       run: |
         sed -i -e "/gxx_linux-64/c\  - gxx_linux-64 ${{ matrix.gcc }}" \
-               -e "/torchani/c\  - torchani ${{ matrix.torchani }}" \
                -e "/python/c\  - python ${{ matrix.python }}" \
                -e "/pytorch-gpu/c\  - pytorch-gpu ${{ matrix.pytorch }}" \
                environment.yml
@@ -91,9 +99,20 @@ jobs:
       run: |
         conda activate nnpops
         mkdir build && cd build
+        SET_TORCH_DIR=""
+        if [ "${{ matrix.set_torch_dir }}" = "true" ]; then
+            # For older PyTorch, this is necessary; for newer PyTorch, CMake can
+            # find it automatically, and including this breaks the build.
+            SET_TORCH_DIR="-DTorch_DIR=$(python -c 'import torch.utils; print(torch.utils.cmake_prefix_path)')/Torch"
+        fi
+        # For newer PyTorch, it is necessary to point CMake to the nvcc that we
+        # installed directly onto the system; otherwise, it will pick up the
+        # one pulled in by conda after installing pytorch-gpu, and break since
+        # the CUDA headers will be missing.
         cmake .. \
               -DENABLE_CUDA=${{ matrix.enable_cuda }} \
-              -DTorch_DIR=$(python -c 'import torch.utils; print(torch.utils.cmake_prefix_path)')/Torch \
+              -DCMAKE_CUDA_COMPILER=/usr/local/cuda/bin/nvcc \
+              $SET_TORCH_DIR \
               -DCMAKE_INSTALL_PREFIX=$CONDA_PREFIX
         make install
 
diff --git a/CMakeLists.txt b/CMakeLists.txt
index a68df575..a2ce279e 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -16,7 +16,6 @@ enable_testing()
 # Source files of the library
 set(SRC_FILES src/ani/CpuANISymmetryFunctions.cpp
               src/ani/CudaANISymmetryFunctions.cu
-              src/pytorch/BatchedNN.cpp
               src/pytorch/CFConv.cpp
               src/pytorch/CFConvNeighbors.cpp
               src/pytorch/SymmetryFunctions.cpp
@@ -75,8 +74,8 @@ add_custom_target(copy_test ALL
 add_custom_command(
   TARGET copy_test POST_BUILD
   COMMAND ${CMAKE_COMMAND} -E copy_directory
-  ${CMAKE_SOURCE_DIR}/src/pytorch/molecules
-  ${CMAKE_BINARY_DIR}/test/molecules
+  ${CMAKE_SOURCE_DIR}/src/pytorch/test_data
+  ${CMAKE_BINARY_DIR}/test/test_data
 )
 
 # Copy each test to the binary dir and add it
@@ -103,12 +102,8 @@ add_test(TestGetNeighborPairs  pytest -v --doctest-modules ${CMAKE_BINARY_DIR}/t
 # Installation
 install(TARGETS ${LIBRARY} DESTINATION ${Python3_SITEARCH}/${NAME})
 install(FILES src/pytorch/__init__.py
-              src/pytorch/BatchedNN.py
               src/pytorch/CFConv.py
               src/pytorch/CFConvNeighbors.py
-              src/pytorch/EnergyShifter.py
-              src/pytorch/OptimizedTorchANI.py
-              src/pytorch/SpeciesConverter.py
               src/pytorch/SymmetryFunctions.py
         DESTINATION ${Python3_SITEARCH}/${NAME})
 install(FILES src/pytorch/neighbors/__init__.py
diff --git a/README.md b/README.md
index 41029c3d..8737e62d 100644
--- a/README.md
+++ b/README.md
@@ -39,7 +39,7 @@ If you don't have `conda`, we recommend installing [Miniconda](https://docs.cond
 #### Prerequisites
 
 - *CUDA Toolkit* (https://developer.nvidia.com/cuda-downloads)
-- *Miniconda* (https://docs.conda.io/en/latest/miniconda.html#linux-installers)
+- *Miniconda* (https://docs.conda.io/en/latest/miniconda.html)
 
 #### Build & install
 
@@ -48,24 +48,20 @@ If you don't have `conda`, we recommend installing [Miniconda](https://docs.cond
 $ git clone https://github.com/openmm/NNPOps.git
 ```
 
-- Set `CUDA_HOME`
-```bash
-$ export CUDA_HOME=/usr/local/cuda-11.2
-```
-
-- Crate and activate a *Conda* environment
+- Create and activate a Conda environment
 ```bash
 $ cd NNPOps
 $ conda env create -n nnpops -f environment.yml
 $ conda activate nnpops
 ```
 
-- Configure, build, and install
+- Configure, build, and install (note: for some older PyTorch versions, you may
+  need to add `-DTorch_DIR=$(python -c 'import torch.utils; print(torch.utils.cmake_prefix_path)')/Torch`
+  as an argument to `cmake`; for newer PyTorch versions, this is unnecessary,
+  and may actually cause CMake configuration to fail)
 ```bash
 $ mkdir build && cd build
-$ cmake .. \
-        -DTorch_DIR=$(python -c 'import torch.utils; print(torch.utils.cmake_prefix_path)')/Torch \
-        -DCMAKE_INSTALL_PREFIX=$CONDA_PREFIX
+$ cmake .. -DCMAKE_INSTALL_PREFIX=$CONDA_PREFIX
 $ make install
 ```
 
@@ -74,56 +70,12 @@ $ make install
 $ ctest --verbose
 ```
 
-## Usage
-
-Accelerated [*TorchANI*](https://aiqm.github.io/torchani/) operations:
-- [`torchani.AEVComputer`](https://aiqm.github.io/torchani/api.html?highlight=speciesaev#torchani.AEVComputer)
-- [`torchani.neurochem.NeuralNetwork`](https://aiqm.github.io/torchani/api.html#module-torchani.neurochem)
-
-### Example
-
-```python
-import mdtraj
-import torch
-import torchani
-
-from NNPOps.SpeciesConverter import TorchANISpeciesConverter
-from NNPOps.SymmetryFunctions import TorchANISymmetryFunctions
-from NNPOps.BatchedNN import TorchANIBatchedNN
-from NNPOps.EnergyShifter import TorchANIEnergyShifter
-
-from NNPOps import OptimizedTorchANI
+## Operations
 
-device = torch.device('cuda')
+The following optimized operations are present in NNPOps and accessible from
+Python using the listed classes or functions:
 
-# Load a molecule
-molecule = mdtraj.load('molecule.mol2')
-species = torch.tensor([[atom.element.atomic_number for atom in molecule.top.atoms]], device=device)
-positions = torch.tensor(molecule.xyz * 10, dtype=torch.float32, requires_grad=True, device=device)
-
-# Construct ANI-2x and replace its operations with the optimized ones
-nnp = torchani.models.ANI2x(periodic_table_index=True).to(device)
-nnp.species_converter = TorchANISpeciesConverter(nnp.species_converter, species).to(device)
-nnp.aev_computer = TorchANISymmetryFunctions(nnp.species_converter, nnp.aev_computer, species).to(device)
-nnp.neural_networks = TorchANIBatchedNN(nnp.species_converter, nnp.neural_networks, species).to(device)
-nnp.energy_shifter = TorchANIEnergyShifter(nnp.species_converter, nnp.energy_shifter, species).to(device)
-
-# Compute energy and forces
-energy = nnp((species, positions)).energies
-energy.backward()
-forces = -positions.grad.clone()
-
-print(energy, forces)
-
-# Alternatively, all the optimizations can be applied with OptimizedTorchANI
-nnp2 = torchani.models.ANI2x(periodic_table_index=True).to(device)
-nnp2 = OptimizedTorchANI(nnp2, species).to(device)
-
-# Compute energy and forces again
-energy = nnp2((species, positions)).energies
-positions.grad.zero_()
-energy.backward()
-forces = -positions.grad.clone()
-
-print(energy, forces)
-```
+- ANI symmetry functions: `NNPOps.SymmetryFunctions.ANISymmetryFunctions`
+- Continuous filter convolution (CFConv): `NNPOps.CFConv.CFConv`, `NNPOps.CFConv.CFConvNeighbors`
+- Neighbor pair enumeration: `NNPOps.neighbors.getNeighborPairs()`
+- Particle mesh Ewald (PME): `NNPOps.pme.PME`
diff --git a/src/pytorch/molecules/1hvj_ligand.mol2 b/devtools/test_data/1hvj_ligand.mol2
similarity index 100%
rename from src/pytorch/molecules/1hvj_ligand.mol2
rename to devtools/test_data/1hvj_ligand.mol2
diff --git a/src/pytorch/molecules/1hvk_ligand.mol2 b/devtools/test_data/1hvk_ligand.mol2
similarity index 100%
rename from src/pytorch/molecules/1hvk_ligand.mol2
rename to devtools/test_data/1hvk_ligand.mol2
diff --git a/src/pytorch/molecules/2iuz_ligand.mol2 b/devtools/test_data/2iuz_ligand.mol2
similarity index 100%
rename from src/pytorch/molecules/2iuz_ligand.mol2
rename to devtools/test_data/2iuz_ligand.mol2
diff --git a/src/pytorch/molecules/3hkw_ligand.mol2 b/devtools/test_data/3hkw_ligand.mol2
similarity index 100%
rename from src/pytorch/molecules/3hkw_ligand.mol2
rename to devtools/test_data/3hkw_ligand.mol2
diff --git a/src/pytorch/molecules/3hky_ligand.mol2 b/devtools/test_data/3hky_ligand.mol2
similarity index 100%
rename from src/pytorch/molecules/3hky_ligand.mol2
rename to devtools/test_data/3hky_ligand.mol2
diff --git a/src/pytorch/molecules/3lka_ligand.mol2 b/devtools/test_data/3lka_ligand.mol2
similarity index 100%
rename from src/pytorch/molecules/3lka_ligand.mol2
rename to devtools/test_data/3lka_ligand.mol2
diff --git a/src/pytorch/molecules/3o99_ligand.mol2 b/devtools/test_data/3o99_ligand.mol2
similarity index 100%
rename from src/pytorch/molecules/3o99_ligand.mol2
rename to devtools/test_data/3o99_ligand.mol2
diff --git a/devtools/test_data/generate.py b/devtools/test_data/generate.py
new file mode 100755
index 00000000..b60fd0a9
--- /dev/null
+++ b/devtools/test_data/generate.py
@@ -0,0 +1,55 @@
+#!/usr/bin/env python
+
+"""
+Uses TorchANI to generate test cases in `src/pytorch/test_data/*.pt` for the
+ANI symmetry functions.  Requires TorchANI and mdtraj to be installed.
+"""
+
+import mdtraj
+import torch
+import torchani
+
+def main():
+    for name in ("1hvj", "1hvk", "2iuz", "3hkw", "3hky", "3lka", "3o99"):
+        generate_molecule_test_case(f"{name}_ligand.mol2", f"{name}.pt")
+    generate_molecule_test_case("water.pdb", "water.pt", True)
+
+def generate_molecule_test_case(in_path, out_path, pbc=False):
+    molecule = mdtraj.load(in_path)
+    atomic_numbers = torch.tensor([atom.element.atomic_number for atom in molecule.top.atoms])
+    atomic_positions = torch.tensor(molecule.xyz[0] * 10, requires_grad=True)
+    cell = torch.tensor(molecule.unitcell_vectors[0] * 10) if pbc else None
+
+    nnp = torchani.models.ANI2x()
+
+    parameters = dict(
+        numSpecies=nnp.aev_computer.num_species,
+        Rcr=nnp.aev_computer.radial.cutoff,
+        Rca=nnp.aev_computer.angular.cutoff,
+        EtaR=nnp.aev_computer.radial.eta.tolist(),
+        ShfR=nnp.aev_computer.radial.shifts.tolist(),
+        EtaA=nnp.aev_computer.angular.eta.tolist(),
+        Zeta=nnp.aev_computer.angular.zeta.tolist(),
+        ShfA=nnp.aev_computer.angular.shifts.tolist(),
+        ShfZ=nnp.aev_computer.angular.sections.tolist(),
+        atomSpecies=nnp.species_converter(atomic_numbers).tolist(),
+    )
+    output = nnp.aev_computer.forward(
+        torch.tensor(parameters["atomSpecies"]).unsqueeze(0),
+        atomic_positions.unsqueeze(0),
+        cell,
+        None if cell is None else torch.tensor([True, True, True])
+    )[0]
+    total = torch.sum(output)
+    total.backward()
+    testcase = dict(
+        parameters=parameters,
+        positions=atomic_positions.detach(),
+        cell=cell,
+        output=output,
+        grad=atomic_positions.grad,
+    )
+    torch.save(testcase, out_path)
+
+if __name__ == "__main__":
+    main()
diff --git a/src/pytorch/molecules/water.pdb b/devtools/test_data/water.pdb
similarity index 100%
rename from src/pytorch/molecules/water.pdb
rename to devtools/test_data/water.pdb
diff --git a/environment.yml b/environment.yml
index b22fdd16..a1ed61a9 100644
--- a/environment.yml
+++ b/environment.yml
@@ -2,12 +2,9 @@ channels:
   - conda-forge
 dependencies:
   - cmake >=3.20
-  - gxx_linux-64 10.3.*
+  - gxx_linux-64
   - make
-  - mdtraj
-  - torchani 2.2.*
   - pytest
-  - python 3.10.*
-  - pytorch-gpu 2.0.*
-  - sysroot_linux-64 2.17
-
+  - python
+  - pytorch-gpu
+  - sysroot_linux-64
diff --git a/src/ani/CudaANISymmetryFunctions.cu b/src/ani/CudaANISymmetryFunctions.cu
index c1ebca86..ffcb75b8 100644
--- a/src/ani/CudaANISymmetryFunctions.cu
+++ b/src/ani/CudaANISymmetryFunctions.cu
@@ -54,6 +54,7 @@ CudaANISymmetryFunctions::CudaANISymmetryFunctions(int numAtoms, int numSpecies,
     CHECK_RESULT(cudaMemcpyAsync(atomSpeciesArray, atomSpecies.data(), atomSpecies.size()*sizeof(int), cudaMemcpyDefault));
     CHECK_RESULT(cudaMemcpyAsync(radialFunctionArray, radialFunctions.data(), radialFunctions.size()*sizeof(RadialFunction), cudaMemcpyDefault));
     CHECK_RESULT(cudaMemcpyAsync(angularFunctionArray, angularFunctions.data(), angularFunctions.size()*sizeof(AngularFunction), cudaMemcpyDefault));
+    CHECK_RESULT(cudaStreamSynchronize(0));
 
     // There are numSpecies*(numSpecies+1)/2 copies of each angular symmetry function.  Create a table mapping from
     // the species indices of two atoms to the corresponding symmetry function index.
@@ -312,7 +313,7 @@ void CudaANISymmetryFunctions::computeSymmetryFunctions(const float* positions,
     float* angularPtr;
     cudaPointerAttributes attrib;
     cudaError_t result = cudaPointerGetAttributes(&attrib, radial);
-    if (result != cudaSuccess || attrib.devicePointer == 0) {
+    if (result != cudaSuccess || attrib.type != cudaMemoryTypeDevice || attrib.devicePointer == 0) {
         radialOnDevice = false;
         radialPtr = radialValues;
     }
@@ -321,7 +322,7 @@ void CudaANISymmetryFunctions::computeSymmetryFunctions(const float* positions,
         radialPtr = (float*) attrib.devicePointer;
     }
     result = cudaPointerGetAttributes(&attrib, angular);
-    if (result != cudaSuccess || attrib.devicePointer == 0) {
+    if (result != cudaSuccess || attrib.type != cudaMemoryTypeDevice || attrib.devicePointer == 0) {
         angularOnDevice = false;
         angularPtr = angularValues;
     }
@@ -332,19 +333,19 @@ void CudaANISymmetryFunctions::computeSymmetryFunctions(const float* positions,
 
     // Record the positions and periodic box vectors.
 
-    CHECK_RESULT(cudaMemcpyAsync(this->positions, positions, 3*numAtoms*sizeof(float), cudaMemcpyDefault));
+    CHECK_RESULT(cudaMemcpyAsync(this->positions, positions, 3*numAtoms*sizeof(float), cudaMemcpyDefault, stream));
     float* hostBoxVectors;
     if (periodic) {
         // We'll need to access the box vectors on both host and device.  Figure out the most
         // efficient way of doing that.
 
         result = cudaPointerGetAttributes(&attrib, periodicBoxVectors);
-        if (result != cudaSuccess || attrib.hostPointer == 0) {
+        if (result != cudaSuccess || attrib.type != cudaMemoryTypeDevice || attrib.hostPointer == 0) {
             CHECK_RESULT(cudaMemcpy(this->periodicBoxVectors, periodicBoxVectors, 9*sizeof(float), cudaMemcpyDefault));
             hostBoxVectors = this->periodicBoxVectors;
         }
         else {
-            CHECK_RESULT(cudaMemcpyAsync(this->periodicBoxVectors, periodicBoxVectors, 9*sizeof(float), cudaMemcpyDefault));
+            CHECK_RESULT(cudaMemcpyAsync(this->periodicBoxVectors, periodicBoxVectors, 9*sizeof(float), cudaMemcpyDefault, stream));
             hostBoxVectors = (float*) attrib.hostPointer;
         }
     }
@@ -360,8 +361,8 @@ void CudaANISymmetryFunctions::computeSymmetryFunctions(const float* positions,
 
     // Clear the output arrays.
 
-    CHECK_RESULT(cudaMemsetAsync(radialPtr, 0, numAtoms*numSpecies*radialFunctions.size()*sizeof(float)));
-    CHECK_RESULT(cudaMemsetAsync(angularPtr, 0, numAtoms*(numSpecies*(numSpecies+1)/2)*angularFunctions.size()*sizeof(float)));
+    CHECK_RESULT(cudaMemsetAsync(radialPtr, 0, numAtoms*numSpecies*radialFunctions.size()*sizeof(float), stream));
+    CHECK_RESULT(cudaMemsetAsync(angularPtr, 0, numAtoms*(numSpecies*(numSpecies+1)/2)*angularFunctions.size()*sizeof(float), stream));
 
     // Compute the symmetry functions.
 
@@ -400,9 +401,9 @@ void CudaANISymmetryFunctions::computeSymmetryFunctions(const float* positions,
     // Copy the final values to the destination memory.
 
     if (!radialOnDevice)
-        CHECK_RESULT(cudaMemcpyAsync(radial, radialValues, numAtoms*numSpecies*radialFunctions.size()*sizeof(float), cudaMemcpyDefault));
+        CHECK_RESULT(cudaMemcpyAsync(radial, radialValues, numAtoms*numSpecies*radialFunctions.size()*sizeof(float), cudaMemcpyDefault, stream));
     if (!angularOnDevice)
-        CHECK_RESULT(cudaMemcpyAsync(angular, angularValues, numAtoms*(numSpecies*(numSpecies+1))*angularFunctions.size()*sizeof(float)/2, cudaMemcpyDefault));
+        CHECK_RESULT(cudaMemcpyAsync(angular, angularValues, numAtoms*(numSpecies*(numSpecies+1))*angularFunctions.size()*sizeof(float)/2, cudaMemcpyDefault, stream));
 }
 
 template <bool PERIODIC, bool TRICLINIC>
@@ -608,21 +609,21 @@ void CudaANISymmetryFunctions::backprop(const float* radialDeriv, const float* a
     float* posPtr;
     cudaPointerAttributes attrib;
     cudaError_t result = cudaPointerGetAttributes(&attrib, radialDeriv);
-    if (result != cudaSuccess || attrib.devicePointer == 0) {
-        CHECK_RESULT(cudaMemcpyAsync(radialValues, radialDeriv, numAtoms*numSpecies*numRadial*sizeof(float), cudaMemcpyDefault));
+    if (result != cudaSuccess || attrib.type != cudaMemoryTypeDevice || attrib.devicePointer == 0) {
+        CHECK_RESULT(cudaMemcpyAsync(radialValues, radialDeriv, numAtoms*numSpecies*numRadial*sizeof(float), cudaMemcpyDefault, stream));
         radialPtr = radialValues;
     }
     else
         radialPtr = (float*) attrib.devicePointer;
     result = cudaPointerGetAttributes(&attrib, angularDeriv);
-    if (result != cudaSuccess || attrib.devicePointer == 0) {
-        CHECK_RESULT(cudaMemcpyAsync(angularValues, angularDeriv, numAtoms*(numSpecies*(numSpecies+1))*numAngular*sizeof(float)/2, cudaMemcpyDefault));
+    if (result != cudaSuccess || attrib.type != cudaMemoryTypeDevice || attrib.devicePointer == 0) {
+        CHECK_RESULT(cudaMemcpyAsync(angularValues, angularDeriv, numAtoms*(numSpecies*(numSpecies+1))*numAngular*sizeof(float)/2, cudaMemcpyDefault, stream));
         angularPtr = angularValues;
     }
     else
         angularPtr = (float*) attrib.devicePointer;
     result = cudaPointerGetAttributes(&attrib, positionDeriv);
-    if (result != cudaSuccess || attrib.devicePointer == 0) {
+    if (result != cudaSuccess || attrib.type != cudaMemoryTypeDevice || attrib.devicePointer == 0) {
         posOnDevice = false;
         posPtr = positionDerivValues;
     }
@@ -633,7 +634,7 @@ void CudaANISymmetryFunctions::backprop(const float* radialDeriv, const float* a
 
     // Clear the output array.
 
-    CHECK_RESULT(cudaMemsetAsync(posPtr, 0, numAtoms*sizeof(float3)));
+    CHECK_RESULT(cudaMemsetAsync(posPtr, 0, numAtoms*sizeof(float3), stream));
 
     // Backpropagate through the symmetry functions.
 
@@ -666,6 +667,6 @@ void CudaANISymmetryFunctions::backprop(const float* radialDeriv, const float* a
     // Copy the final values to the destination memory.
 
     if (!posOnDevice)
-        CHECK_RESULT(cudaMemcpyAsync(positionDeriv, positionDerivValues, numAtoms*sizeof(float3), cudaMemcpyDefault));
+        CHECK_RESULT(cudaMemcpyAsync(positionDeriv, positionDerivValues, numAtoms*sizeof(float3), cudaMemcpyDefault, stream));
 }
 
diff --git a/src/pytorch/BatchedNN.cpp b/src/pytorch/BatchedNN.cpp
deleted file mode 100644
index 75b42e1c..00000000
--- a/src/pytorch/BatchedNN.cpp
+++ /dev/null
@@ -1,50 +0,0 @@
-/**
- * Copyright (c) 2020 Acellera
- * Authors: Raimondas Galvelis
- * 
- * Permission is hereby granted, free of charge, to any person obtaining a copy
- * of this software and associated documentation files (the "Software"), to deal
- * in the Software without restriction, including without limitation the rights
- * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
- * copies of the Software, and to permit persons to whom the Software is
- * furnished to do so, subject to the following conditions:
- * 
- * The above copyright notice and this permission notice shall be included in all
- * copies or substantial portions of the Software.
- * 
- * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
- * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
- * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
- * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
- * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
- * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
- * SOFTWARE.
- */
-
-#include <torch/script.h>
-
-using Context = torch::autograd::AutogradContext;
-using Tensor = torch::Tensor;
-using tensor_list = torch::autograd::tensor_list;
-
-class BatchedLinearFunction : public torch::autograd::Function<BatchedLinearFunction> {
-public:
-    static Tensor forward(Context* ctx, const Tensor& vectors, const Tensor& weights, const Tensor& biases) {
-        ctx->save_for_backward({weights});
-        return torch::matmul(weights, vectors) + biases;
-    };
-    static tensor_list backward(Context *ctx, const tensor_list& grads) {
-        const Tensor grad_in = grads[0].squeeze(-1).unsqueeze(-2);
-        const Tensor weights = ctx->get_saved_variables()[0];
-        const Tensor grad_out = torch::matmul(grad_in, weights).squeeze(-2).unsqueeze(-1);
-        return {grad_out, torch::Tensor(), torch::Tensor()};
-    };
-};
-
-static Tensor BatchedLinear(const Tensor& vector, const Tensor& weights, const Tensor& biases) {
-    return BatchedLinearFunction::apply(vector, weights, biases);
-}
-
-TORCH_LIBRARY(NNPOpsBatchedNN, m) {
-    m.def("BatchedLinear", BatchedLinear);
-}
\ No newline at end of file
diff --git a/src/pytorch/BatchedNN.py b/src/pytorch/BatchedNN.py
deleted file mode 100644
index 79f51ac0..00000000
--- a/src/pytorch/BatchedNN.py
+++ /dev/null
@@ -1,122 +0,0 @@
-#
-# Copyright (c) 2020 Acellera
-# Authors: Raimondas Galvelis
-#
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-#
-# The above copyright notice and this permission notice shall be included in all
-# copies or substantial portions of the Software.
-#
-# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-# SOFTWARE.
-#
-
-import torch
-from torch import nn
-from torch import Tensor
-from torch.nn import functional as F
-from typing import List, NamedTuple, Tuple, Union
-
-batchedLinear = torch.ops.NNPOpsBatchedNN.BatchedLinear
-
-
-class SpeciesEnergies(NamedTuple):
-    species: Tensor
-    energies: Tensor
-
-class _BatchedNN(torch.nn.Module):
-
-    from torchani.nn import ANIModel, Ensemble, SpeciesConverter # https://github.com/openmm/NNPOps/issues/44
-
-    def __init__(self, converter: SpeciesConverter, ensemble: Union[ANIModel, Ensemble], atomicNumbers: Tensor):
-
-        super().__init__()
-
-        # Convert atomic numbers to a list of species
-        species_list = converter((atomicNumbers, torch.empty(0))).species[0].tolist()
-
-        # Handle the case when the ensemble is just one model
-        self._ensemble = ensemble if isinstance(ensemble, torch.nn.ModuleList) else [ensemble]
-
-        # Convert models to the list of linear layers
-        models = [list(model.values()) for model in self._ensemble]
-
-        # Extract the weihts and biases of the linear layers
-        for ilayer in [0, 2, 4, 6]:
-            layers = [[model[species][ilayer] for species in species_list] for model in models]
-            weights, biases = self.batchLinearLayers(layers)
-            self.register_buffer(f'layer{ilayer}_weights', weights)
-            self.register_buffer(f'layer{ilayer}_biases', biases)
-
-        # Disable autograd for the parameters
-        for parameter in self.parameters():
-            parameter.requires_grad = False
-
-    @staticmethod
-    def batchLinearLayers(layers: List[List[nn.Linear]]) -> Tuple[Tensor, Tensor]:
-
-        num_models = len(layers)
-        num_atoms = len(layers[0])
-
-        # Note: different elements have different size linear layers, so we just find maximum sizes
-        #       and pad with zeros.
-        max_out = max(layer.out_features for layer in sum(layers, []))
-        max_in = max(layer.in_features for layer in sum(layers, []))
-
-        # Copy weights and biases
-        weights = torch.zeros((1, num_atoms, num_models, max_out, max_in), dtype=torch.float32)
-        biases  = torch.zeros((1, num_atoms, num_models, max_out,      1), dtype=torch.float32)
-        for imodel, sublayers in enumerate(layers):
-            for iatom, layer in enumerate(sublayers):
-                num_out, num_in = layer.weight.shape
-                weights[0, iatom, imodel, :num_out, :num_in] = layer.weight
-                biases [0, iatom, imodel, :num_out,       0] = layer.bias
-
-        return weights, biases
-
-    def _atomic_energies(self, species_aev: Tuple[Tensor, Tensor]) -> Tensor:
-        return self._ensemble[0]._atomic_energies(species_aev)
-
-    def forward(self, species_aev: Tuple[Tensor, Tensor]) -> SpeciesEnergies:
-
-        species, aev = species_aev
-
-        # Reshape: [num_mols, num_atoms, num_features] --> [num_mols, num_atoms, 1, num_features, 1]
-        vectors = aev.unsqueeze(-2).unsqueeze(-1)
-
-        vectors = batchedLinear(vectors, self.layer0_weights, self.layer0_biases) # Linear 0
-        vectors = F.celu(vectors, alpha=0.1)                                      # CELU   1
-        vectors = batchedLinear(vectors, self.layer2_weights, self.layer2_biases) # Linear 2
-        vectors = F.celu(vectors, alpha=0.1)                                      # CELU   3
-        vectors = batchedLinear(vectors, self.layer4_weights, self.layer4_biases) # Linear 4
-        vectors = F.celu(vectors, alpha=0.1)                                      # CELU   5
-        vectors = batchedLinear(vectors, self.layer6_weights, self.layer6_biases) # Linear 6
-
-        # Sum: [num_mols, num_atoms, num_models, 1, 1] --> [num_mols, num_models]
-        # Mean: [num_mols, num_models] --> [num_mols]
-        # The sum and mean must be combined into a single operation to avoid a PyTorch bug.
-        # See https://github.com/openmm/openmm/issues/3812.
-        energies = torch.sum(vectors, (1, 2, 3, 4))/vectors.shape[2]
-
-        return SpeciesEnergies(species, energies)
-
-
-class TorchANIBatchedNN(torch.nn.ModuleList):
-
-    from torchani.nn import ANIModel, Ensemble, SpeciesConverter # https://github.com/openmm/NNPOps/issues/44
-
-    def __init__(self, converter: SpeciesConverter, ensemble: Union[ANIModel, Ensemble], atomicNumbers: Tensor):
-        super().__init__([_BatchedNN(converter, ensemble, atomicNumbers)])
-
-    def forward(self, species_aev: Tuple[Tensor, Tensor]) -> SpeciesEnergies:
-        return self[0].forward(species_aev)
diff --git a/src/pytorch/BenchmarkBatchedNN.py b/src/pytorch/BenchmarkBatchedNN.py
deleted file mode 100644
index b5126d46..00000000
--- a/src/pytorch/BenchmarkBatchedNN.py
+++ /dev/null
@@ -1,99 +0,0 @@
-#
-# Copyright (c) 2020 Acellera
-# Authors: Raimondas Galvelis
-#
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-#
-# The above copyright notice and this permission notice shall be included in all
-# copies or substantial portions of the Software.
-#
-# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-# SOFTWARE.
-#
-
-import mdtraj
-import time
-import torch
-import torchani
-
-# from NNPOps.SymmetryFunctions import TorchANISymmetryFunctions
-from NNPOps.BatchedNN import TorchANIBatchedNN
-
-device = torch.device('cuda')
-
-mol = mdtraj.load('molecules/2iuz_ligand.mol2')
-species = torch.tensor([[atom.element.atomic_number for atom in mol.top.atoms]], device=device)
-positions = torch.tensor(mol.xyz, dtype=torch.float32, requires_grad=True, device=device)
-
-nnp = torchani.models.ANI2x(periodic_table_index=True, model_index=None).to(device)
-print(nnp)
-
-energy_ref = nnp((species, positions)).energies
-energy_ref.backward()
-grad_ref = positions.grad.clone()
-
-N = 3000
-start = time.time()
-for _ in range(N):
-    energy_ref = nnp((species, positions)).energies
-delta = time.time() - start
-print(f'ANI-2x (forward pass)')
-print(f'  Duration: {delta} s')
-print(f'  Speed: {delta/N*1000} ms/it')
-
-N = 1000
-start = time.time()
-for _ in range(N):
-    energy_ref = nnp((species, positions)).energies
-    positions.grad.zero_()
-    energy_ref.backward()
-delta = time.time() - start
-print(f'ANI-2x (forward & backward pass)')
-print(f'  Duration: {delta} s')
-print(f'  Speed: {delta/N*1000} ms/it')
-
-# nnp.aev_computer = TorchANISymmetryFunctions(nnp.aev_computer).to(device)
-nnp.neural_networks = TorchANIBatchedNN(nnp.species_converter, nnp.neural_networks, species).to(device)
-print(nnp)
-
-# nnp = torch.jit.script(nnp)
-# nnp.save('nnp.pt')
-# npp = torch.jit.load('nnp.pt').to(device)
-
-energy = nnp((species, positions)).energies
-positions.grad.zero_()
-energy.backward()
-grad = positions.grad.clone()
-
-N = 15000
-start = time.time()
-for _ in range(N):
-    energy = nnp((species, positions)).energies
-delta = time.time() - start
-print(f'ANI-2x with BatchedNN (forward pass)')
-print(f'  Duration: {delta} s')
-print(f'  Speed: {delta/N*1000} ms/it')
-
-N = 7500
-start = time.time()
-for _ in range(N):
-    energy = nnp((species, positions)).energies
-    positions.grad.zero_()
-    energy.backward()
-delta = time.time() - start
-print(f'ANI-2x with BatchedNN (forward & backward pass)')
-print(f'  Duration: {delta} s')
-print(f'  Speed: {delta/N*1000} ms/it')
-
-# print(float(energy_ref), float(energy), float(energy_ref - energy))
-# print(float(torch.max(torch.abs((grad - grad_ref)/grad_ref))))
\ No newline at end of file
diff --git a/src/pytorch/BenchmarkTorchANISymmetryFunctions.py b/src/pytorch/BenchmarkTorchANISymmetryFunctions.py
deleted file mode 100644
index 6299c73e..00000000
--- a/src/pytorch/BenchmarkTorchANISymmetryFunctions.py
+++ /dev/null
@@ -1,59 +0,0 @@
-import mdtraj
-import time
-import torch
-import torchani
-
-from NNPOps.SymmetryFunctions import TorchANISymmetryFunctions
-
-device = torch.device('cuda')
-
-mol = mdtraj.load('molecules/2iuz_ligand.mol2')
-species = torch.tensor([[atom.element.atomic_number for atom in mol.top.atoms]], device=device)
-positions = torch.tensor(mol.xyz * 10, dtype=torch.float32, requires_grad=True, device=device)
-
-nnp = torchani.models.ANI2x(periodic_table_index=True, model_index=None).to(device)
-speciesPositions = nnp.species_converter((species, positions))
-symmFuncRef = nnp.aev_computer
-symmFunc = TorchANISymmetryFunctions(nnp.aev_computer).to(device)
-
-aev_ref = symmFuncRef(speciesPositions).aevs
-sum_aev_ref = torch.sum(aev_ref)
-sum_aev_ref.backward()
-grad_ref = positions.grad.clone()
-
-N = 10000
-start = time.time()
-for _ in range(N):
-    aev_ref = symmFuncRef(speciesPositions).aevs
-    sum_aev_ref = torch.sum(aev_ref)
-    positions.grad.zero_()
-    sum_aev_ref.backward()
-delta = time.time() - start
-grad_ref = positions.grad.clone()
-print('Original TorchANI symmetry functions')
-print(f'  Duration: {delta} s')
-print(f'  Speed: {delta/N*1000} ms/it')
-
-aev = symmFunc(speciesPositions).aevs
-sum_aev = torch.sum(aev)
-positions.grad.zero_()
-sum_aev.backward()
-grad = positions.grad.clone()
-
-N = 100000
-start = time.time()
-for _ in range(N):
-    aev = symmFunc(speciesPositions).aevs
-    sum_aev = torch.sum(aev)
-    positions.grad.zero_()
-    sum_aev.backward()
-delta = time.time() - start
-grad = positions.grad.clone()
-print('Optimized TorchANI symmetry functions')
-print(f'  Duration: {delta} s')
-print(f'  Speed: {delta/N*1000} ms/it')
-
-aev_error = torch.max(torch.abs(aev - aev_ref))
-grad_error = torch.max(torch.abs(grad - grad_ref))
-assert aev_error < 0.0002
-assert grad_error < 0.007
\ No newline at end of file
diff --git a/src/pytorch/EnergyShifter.py b/src/pytorch/EnergyShifter.py
deleted file mode 100644
index 75e289ab..00000000
--- a/src/pytorch/EnergyShifter.py
+++ /dev/null
@@ -1,53 +0,0 @@
-#
-# Copyright (c) 2020-2021 Acellera
-# Authors: Raimondas Galvelis
-#
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-#
-# The above copyright notice and this permission notice shall be included in all
-# copies or substantial portions of the Software.
-#
-# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-# SOFTWARE.
-#
-
-import torch
-from torch import Tensor
-from typing import NamedTuple, Optional, Tuple
-
-class SpeciesEnergies(NamedTuple):
-    species: Tensor
-    energies: Tensor
-
-class TorchANIEnergyShifter(torch.nn.Module):
-
-    from torchani.nn import SpeciesConverter # https://github.com/openmm/NNPOps/issues/44
-    from torchani.utils import EnergyShifter # https://github.com/openmm/NNPOps/issues/44
-
-    def __init__(self, converter: SpeciesConverter, shifter: EnergyShifter, atomicNumbers: Tensor) -> None:
-
-        super().__init__()
-
-        # Convert atomic numbers to a list of species
-        species = converter((atomicNumbers, torch.empty(0))).species
-
-        # Compute atomic self energies
-        self.register_buffer('self_energies', shifter.sae(species))
-
-    def forward(self, species_energies: Tuple[Tensor, Tensor],
-                cell: Optional[Tensor] = None,
-                pbc: Optional[Tensor] = None) -> SpeciesEnergies:
-
-        species, energies = species_energies
-
-        return SpeciesEnergies(species, energies + self.self_energies)
\ No newline at end of file
diff --git a/src/pytorch/OptimizedTorchANI.py b/src/pytorch/OptimizedTorchANI.py
deleted file mode 100644
index 96819c21..00000000
--- a/src/pytorch/OptimizedTorchANI.py
+++ /dev/null
@@ -1,54 +0,0 @@
-#
-# Copyright (c) 2020-2021 Acellera
-# Authors: Raimondas Galvelis
-#
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-#
-# The above copyright notice and this permission notice shall be included in all
-# copies or substantial portions of the Software.
-#
-# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-# SOFTWARE.
-#
-
-import torch
-from torch import Tensor
-from typing import Optional, Tuple
-
-from NNPOps.BatchedNN import TorchANIBatchedNN
-from NNPOps.EnergyShifter import TorchANIEnergyShifter, SpeciesEnergies
-from NNPOps.SpeciesConverter import TorchANISpeciesConverter
-from NNPOps.SymmetryFunctions import TorchANISymmetryFunctions
-
-class OptimizedTorchANI(torch.nn.Module):
-
-    def __init__(self, model, atomicNumbers: Tensor) -> None:
-
-        super().__init__()
-
-        # Optimize the components of an ANI model
-        self.species_converter = TorchANISpeciesConverter(model.species_converter, atomicNumbers)
-        self.aev_computer = TorchANISymmetryFunctions(model.species_converter, model.aev_computer, atomicNumbers)
-        self.neural_networks = TorchANIBatchedNN(model.species_converter, model.neural_networks, atomicNumbers)
-        self.energy_shifter = TorchANIEnergyShifter(model.species_converter, model.energy_shifter, atomicNumbers)
-
-    def forward(self, species_coordinates: Tuple[Tensor, Tensor],
-                cell: Optional[Tensor] = None,
-                pbc: Optional[Tensor] = None) -> SpeciesEnergies:
-
-        species_coordinates = self.species_converter(species_coordinates)
-        species_aevs = self.aev_computer(species_coordinates, cell, pbc)
-        species_energies = self.neural_networks(species_aevs)
-        species_energies = self.energy_shifter(species_energies)
-
-        return species_energies
diff --git a/src/pytorch/SpeciesConverter.py b/src/pytorch/SpeciesConverter.py
deleted file mode 100644
index 0bb6990d..00000000
--- a/src/pytorch/SpeciesConverter.py
+++ /dev/null
@@ -1,44 +0,0 @@
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-#
-# The above copyright notice and this permission notice shall be included in all
-# copies or substantial portions of the Software.
-#
-# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-# SOFTWARE.
-#
-
-import torch
-from torch import Tensor
-from typing import NamedTuple, Optional, Tuple
-
-class SpeciesCoordinates(NamedTuple):
-    species: Tensor
-    coordinates: Tensor
-
-class TorchANISpeciesConverter(torch.nn.Module):
-
-    from torchani.nn import SpeciesConverter
-
-    def __init__(self, converter: SpeciesConverter, atomicNumbers: Tensor) -> None:
-
-        super().__init__()
-
-        # Convert atomic numbers to a list of species
-        species = converter((atomicNumbers, torch.empty(0))).species
-        self.register_buffer('species', species)
-
-        self.conv_tensor = converter.conv_tensor # Just to make TorchScript happy :)
-
-    def forward(self, species_coordinates: Tuple[Tensor, Tensor],
-                cell: Optional[Tensor] = None,
-                pbc: Optional[Tensor] = None) -> Tuple[Tensor, Tensor]:
-
-        _, coordinates = species_coordinates
-
-        return SpeciesCoordinates(self.species, coordinates)
diff --git a/src/pytorch/SymmetryFunctions.py b/src/pytorch/SymmetryFunctions.py
index 99dd78ed..f22e5ec8 100644
--- a/src/pytorch/SymmetryFunctions.py
+++ b/src/pytorch/SymmetryFunctions.py
@@ -1,6 +1,6 @@
 #
-# Copyright (c) 2020 Acellera
-# Authors: Raimondas Galvelis
+# Copyright (c) 2020 Acellera, 2025 Stanford University and the Authors
+# Authors: Raimondas Galvelis, Evan Pretti
 #
 # Permission is hereby granted, free of charge, to any person obtaining a copy
 # of this software and associated documentation files (the "Software"), to deal
@@ -21,103 +21,68 @@
 # SOFTWARE.
 #
 
-from typing import List, Optional, Tuple
 import torch
 from torch import Tensor
+from typing import Optional
 
-Holder = torch.classes.NNPOpsANISymmetryFunctions.Holder
-operation = torch.ops.NNPOpsANISymmetryFunctions.operation
-
-class TorchANISymmetryFunctions(torch.nn.Module):
-    """Optimized TorchANI symmetry functions
-
-    Optimized drop-in replacement for torchani.AEVComputer (https://aiqm.github.io/torchani/api.html?highlight=speciesaev#torchani.AEVComputer)
-
-    Example::
-
-        >>> import mdtraj
-        >>> import torch
-        >>> import torchani
-
-        >>> from NNPOps.SymmetryFunctions import TorchANISymmetryFunctions
-
-        >>> device = torch.device('cuda')
-
-        # Load a molecule
-        >>> molecule = mdtraj.load('molecule.mol2')
-        >>> species = torch.tensor([[atom.element.atomic_number for atom in molecule.top.atoms]], device=device)
-        >>> positions = torch.tensor(molecule.xyz * 10, dtype=torch.float32, requires_grad=True, device=device)
-
-        # Construct ANI-2x and replace its native featurizer with NNPOps implementation
-        >>> nnp = torchani.models.ANI2x(periodic_table_index=True).to(device)
-        >>> nnp.aev_computer = TorchANISymmetryFunctions(nnp.species_converter, nnp.aev_computer, species)
-
-        # Compute energy
-        >>> energy = nnp((species, positions)).energies
-        >>> energy.backward()
-        >>> forces = -positions.grad.clone()
-
-        >>> print(energy, forces)
+class ANISymmetryFunctions(torch.nn.Module):
+    """
+    PyTorch module for optimized ANI symmetry functions.
     """
 
-    from torchani import AEVComputer # https://github.com/openmm/NNPOps/pull/38
-    from torchani import SpeciesConverter # https://github.com/openmm/NNPOps/pull/38
+    Holder = torch.classes.NNPOpsANISymmetryFunctions.Holder
+    operation = torch.ops.NNPOpsANISymmetryFunctions.operation
 
-    def __init__(self, converter: SpeciesConverter, symmFunc: AEVComputer, atomicNumbers: Tensor) -> None:
+    def __init__(self, numSpecies: int, Rcr: float, Rca: float,
+                 EtaR: list[float], ShfR: list[float], EtaA: list[float],
+                 Zeta: list[float], ShfA: list[float], ShfZ: list[float],
+                 atomSpecies: list[int]) -> None:
         """
-        Arguments:
-            converter: an instance of torchani.nn.SpeciesConverter (https://aiqm.github.io/torchani/api.html#torchani.SpeciesConverter)
-            symmFunc: an instance of torchani.AEVComputer (https://aiqm.github.io/torchani/api.html#torchani.AEVComputer)
-            atomicNumbers: a tesnor of atomic numbers, e.g. [[6, 1, ,1 ,1, 1]]
+        Create an `ANISymmetryFunctions` instance.
+
+        Parameters
+        ----------
+        numSpecies : int
+            The number of species.
+        Rcr : float
+            The cutoff distance for the radial symmetry functions.
+        Rca : float
+            The cutoff distance for the angular symmetry functions.
+        EtaR : list[float]
+            The Gaussian scale parameters for the radial symmetry functions.
+        ShfR : list[float]
+            The Gaussian shift parameters for the radial symmetry functions.
+        EtaA : list[float]
+            The Gaussian scale parameters for the angular symmetry functions.
+        Zeta : list[float]
+            The exponents for the angular symmetry functions.
+        ShfA : list[float]
+            The Gaussian shift parameters for the angular symmetry functions.
+        ShfZ : list[float]
+            The shift angles for the angular symmetry functions.
+        atomSpecies: list[int]
+            The species indices for each of the atoms.
         """
-        super().__init__()
-
-        self.num_species = symmFunc.num_species
-        Rcr = symmFunc.Rcr
-        Rca = symmFunc.Rca
-        EtaR = symmFunc.EtaR[:, 0].tolist()
-        ShfR = symmFunc.ShfR[0, :].tolist()
-        EtaA = symmFunc.EtaA[:, 0, 0, 0].tolist()
-        Zeta = symmFunc.Zeta[0, :, 0, 0].tolist()
-        ShfA = symmFunc.ShfA[0, 0, :, 0].tolist()
-        ShfZ = symmFunc.ShfZ[0, 0, 0, :].tolist()
-
-        # Convert atomic numbers to species
-        species = converter((atomicNumbers, torch.empty(0))).species[0].tolist()
-
-        # Create a holder
-        self.holder = Holder(self.num_species, Rcr, Rca, EtaR, ShfR, EtaA, Zeta, ShfA, ShfZ, species)
-
-        self.triu_index = torch.tensor([0]) # A dummy variable to make TorchScript happy ;)
 
-    def forward(self, species_positions: Tuple[Tensor, Tensor],
-                      cell: Optional[Tensor] = None,
-                      pbc: Optional[Tensor] = None) -> Tuple[Tensor, Tensor]:
-        """Compute the atomic environment vectors
-
-        The signature of the method is identical to torchani.AEVComputer.forward (https://aiqm.github.io/torchani/api.html?highlight=speciesaev#torchani.AEVComputer.forward)
-
-        Arguments:
-            species_positions: atomic species and positions
-            cell: unitcell vectors
-            pbc: periodic boundary conditions
+        super().__init__()
 
-        Returns:
-            SpeciesAEV: atomic species and environment vectors
+        self.holder = ANISymmetryFunctions.Holder(numSpecies, Rcr, Rca, EtaR, ShfR, EtaA, Zeta, ShfA, ShfZ, atomSpecies)
 
+    def forward(self, positions: Tensor, cell: Optional[Tensor] = None) -> list[Tensor]:
+        """
+        Evaluate the ANI symmetry functions.
+
+        Parameters
+        ----------
+        positions : Tensor
+            Atomic positions.
+        cell : Tensor, optional
+            Box vectors for periodic boundary conditions, if provided.
+
+        Returns
+        -------
+        [Tensor, Tensor]
+            Values of the radial and angular symmetry functions for each atom.
         """
-        species, positions = species_positions
-        if species.shape[0] != 1:
-            raise ValueError('Batched computation of molecules is not supported')
-        if cell is not None:
-            if pbc is None:
-                raise ValueError('"pbc" has to be defined')
-            else:
-                pbc_: List[bool] = pbc.tolist() # Explicit type casting for TorchScript
-                if pbc_ != [True, True, True]:
-                    raise ValueError('Only fully periodic systems are supported, i.e. pbc = [True, True, True]')
-
-        radial, angular = operation(self.holder, positions[0], cell)
-        features = torch.cat((radial, angular), dim=1).unsqueeze(0)
 
-        return species, features
+        return ANISymmetryFunctions.operation(self.holder, positions, cell)
diff --git a/src/pytorch/TestBatchedNN.py b/src/pytorch/TestBatchedNN.py
deleted file mode 100644
index bf0e4c64..00000000
--- a/src/pytorch/TestBatchedNN.py
+++ /dev/null
@@ -1,123 +0,0 @@
-#
-# Copyright (c) 2020 Acellera
-# Authors: Raimondas Galvelis
-#
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-#
-# The above copyright notice and this permission notice shall be included in all
-# copies or substantial portions of the Software.
-#
-# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-# SOFTWARE.
-#
-
-import mdtraj
-import os
-import pytest
-import tempfile
-import torch
-import torchani
-
-molecules = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'molecules')
-
-def test_import():
-    import NNPOps
-    import NNPOps.BatchedNN
-
-class DeterministicTorch:
-    def __enter__(self):
-        if torch.are_deterministic_algorithms_enabled():
-            self._already_enabled = True
-            return
-        self._already_enabled = False
-        torch.use_deterministic_algorithms(True)
-
-    def __exit__(self, type, value, traceback):
-        if not self._already_enabled:
-            torch.use_deterministic_algorithms(False)
-
-@pytest.mark.parametrize('deviceString', ['cpu', 'cuda'])
-@pytest.mark.parametrize('molFile', ['1hvj', '1hvk', '2iuz', '3hkw', '3hky', '3lka', '3o99'])
-def test_compare_with_native(deviceString, molFile):
-
-    if deviceString == 'cuda' and not torch.cuda.is_available():
-        pytest.skip('CUDA is not available')
-    with DeterministicTorch():
-        from NNPOps.BatchedNN import TorchANIBatchedNN
-
-        device = torch.device(deviceString)
-
-        mol = mdtraj.load(os.path.join(molecules, f'{molFile}_ligand.mol2'))
-        atomicNumbers = torch.tensor([[atom.element.atomic_number for atom in mol.top.atoms]], device=device)
-        atomicPositions = torch.tensor(mol.xyz, dtype=torch.float32, requires_grad=True, device=device)
-
-        nnp = torchani.models.ANI2x(periodic_table_index=True).to(device)
-        energy_ref = nnp((atomicNumbers, atomicPositions)).energies
-        energy_ref.backward()
-        grad_ref = atomicPositions.grad.clone()
-
-        nnp.neural_networks = TorchANIBatchedNN(nnp.species_converter, nnp.neural_networks, atomicNumbers).to(device)
-        energy = nnp((atomicNumbers, atomicPositions)).energies
-        atomicPositions.grad.zero_()
-        energy.backward()
-        grad = atomicPositions.grad.clone()
-
-        energy_error = torch.abs((energy - energy_ref)/energy_ref)
-        grad_error = torch.max(torch.abs((grad - grad_ref)/grad_ref))
-
-        assert energy_error < 5e-7
-        if molFile == '3o99':
-            assert grad_error < 0.025 # Some numerical instability
-        else:
-            assert grad_error < 5e-3
-
-@pytest.mark.parametrize('deviceString', ['cpu', 'cuda'])
-@pytest.mark.parametrize('molFile', ['1hvj', '1hvk', '2iuz', '3hkw', '3hky', '3lka', '3o99'])
-def test_model_serialization(deviceString, molFile):
-
-    if deviceString == 'cuda' and not torch.cuda.is_available():
-        pytest.skip('CUDA is not available')
-
-    from NNPOps.BatchedNN import TorchANIBatchedNN
-
-    device = torch.device(deviceString)
-
-    mol = mdtraj.load(os.path.join(molecules, f'{molFile}_ligand.mol2'))
-    atomicNumbers = torch.tensor([[atom.element.atomic_number for atom in mol.top.atoms]], device=device)
-    atomicPositions = torch.tensor(mol.xyz, dtype=torch.float32, requires_grad=True, device=device)
-
-    nnp_ref = torchani.models.ANI2x(periodic_table_index=True).to(device)
-    nnp_ref.neural_networks = TorchANIBatchedNN(nnp_ref.species_converter, nnp_ref.neural_networks, atomicNumbers).to(device)
-
-    energy_ref = nnp_ref((atomicNumbers, atomicPositions)).energies
-    energy_ref.backward()
-    grad_ref = atomicPositions.grad.clone()
-
-    with tempfile.NamedTemporaryFile() as fd:
-
-        torch.jit.script(nnp_ref).save(fd.name)
-        nnp = torch.jit.load(fd.name)
-
-        energy = nnp((atomicNumbers, atomicPositions)).energies
-        atomicPositions.grad.zero_()
-        energy.backward()
-        grad = atomicPositions.grad.clone()
-
-    energy_error = torch.abs((energy - energy_ref)/energy_ref)
-    grad_error = torch.max(torch.abs((grad - grad_ref)/grad_ref))
-
-    assert energy_error < 5e-7
-    if molFile == '3o99':
-        assert grad_error < 0.05 # Some numerical instability
-    else:
-        assert grad_error < 5e-3
diff --git a/src/pytorch/TestCFConv.py b/src/pytorch/TestCFConv.py
index fbbe1863..d658e363 100644
--- a/src/pytorch/TestCFConv.py
+++ b/src/pytorch/TestCFConv.py
@@ -26,8 +26,6 @@
 import tempfile
 import torch
 
-molecules = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'molecules')
-
 def getCFConv(numFilters, device):
 
     from NNPOps.CFConvNeighbors import CFConvNeighbors
diff --git a/src/pytorch/TestCFConvNeighbors.py b/src/pytorch/TestCFConvNeighbors.py
index 27d08e83..3092a53d 100644
--- a/src/pytorch/TestCFConvNeighbors.py
+++ b/src/pytorch/TestCFConvNeighbors.py
@@ -26,8 +26,6 @@
 import tempfile
 import torch
 
-molecules = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'molecules')
-
 def test_import():
     import NNPOps
     import NNPOps.CFConvNeighbors
diff --git a/src/pytorch/TestEnergyShifter.py b/src/pytorch/TestEnergyShifter.py
deleted file mode 100644
index 77e28a07..00000000
--- a/src/pytorch/TestEnergyShifter.py
+++ /dev/null
@@ -1,105 +0,0 @@
-#
-# Copyright (c) 2020-2021 Acellera
-# Authors: Raimondas Galvelis
-#
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-#
-# The above copyright notice and this permission notice shall be included in all
-# copies or substantial portions of the Software.
-#
-# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-# SOFTWARE.
-#
-
-import mdtraj
-import os
-import pytest
-import tempfile
-import torch
-import torchani
-
-molecules = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'molecules')
-
-def test_import():
-    import NNPOps
-    import NNPOps.EnergyShifter
-
-@pytest.mark.parametrize('deviceString', ['cpu', 'cuda'])
-@pytest.mark.parametrize('molFile', ['1hvj', '1hvk', '2iuz', '3hkw', '3hky', '3lka', '3o99'])
-def test_compare_with_native(deviceString, molFile):
-
-    if deviceString == 'cuda' and not torch.cuda.is_available():
-        pytest.skip('CUDA is not available')
-
-    from NNPOps.EnergyShifter import TorchANIEnergyShifter
-
-    device = torch.device(deviceString)
-
-    mol = mdtraj.load(os.path.join(molecules, f'{molFile}_ligand.mol2'))
-    atomicNumbers = torch.tensor([[atom.element.atomic_number for atom in mol.top.atoms]], device=device)
-    atomicPositions = torch.tensor(mol.xyz * 10, dtype=torch.float32, requires_grad=True, device=device)
-
-    nnp = torchani.models.ANI2x(periodic_table_index=True).to(device)
-    energy_ref = nnp((atomicNumbers, atomicPositions)).energies
-    energy_ref.backward()
-    grad_ref = atomicPositions.grad.clone()
-
-    nnp.energy_shifter = TorchANIEnergyShifter(nnp.species_converter, nnp.energy_shifter, atomicNumbers).to(device)
-    energy = nnp((atomicNumbers, atomicPositions)).energies
-    atomicPositions.grad.zero_()
-    energy.backward()
-    grad = atomicPositions.grad.clone()
-
-    energy_error = torch.abs((energy - energy_ref)/energy_ref)
-    grad_error = torch.max(torch.abs((grad - grad_ref)/grad_ref))
-
-    assert energy_error < 5e-7
-    assert grad_error < 5e-3
-
-@pytest.mark.parametrize('deviceString', ['cpu', 'cuda'])
-@pytest.mark.parametrize('molFile', ['1hvj', '1hvk', '2iuz', '3hkw', '3hky', '3lka', '3o99'])
-def test_model_serialization(deviceString, molFile):
-
-    if deviceString == 'cuda' and not torch.cuda.is_available():
-        pytest.skip('CUDA is not available')
-
-    from NNPOps.EnergyShifter import TorchANIEnergyShifter
-
-    device = torch.device(deviceString)
-
-    mol = mdtraj.load(os.path.join(molecules, f'{molFile}_ligand.mol2'))
-    atomicNumbers = torch.tensor([[atom.element.atomic_number for atom in mol.top.atoms]], device=device)
-    atomicPositions = torch.tensor(mol.xyz * 10, dtype=torch.float32, requires_grad=True, device=device)
-
-    nnp_ref = torchani.models.ANI2x(periodic_table_index=True).to(device)
-    nnp_ref.energy_shifter = TorchANIEnergyShifter(nnp_ref.species_converter, nnp_ref.energy_shifter, atomicNumbers).to(device)
-
-    energy_ref = nnp_ref((atomicNumbers, atomicPositions)).energies
-    energy_ref.backward()
-    grad_ref = atomicPositions.grad.clone()
-
-    with tempfile.NamedTemporaryFile() as fd:
-
-        torch.jit.script(nnp_ref).save(fd.name)
-        nnp = torch.jit.load(fd.name)
-
-        energy = nnp((atomicNumbers, atomicPositions)).energies
-        atomicPositions.grad.zero_()
-        energy.backward()
-        grad = atomicPositions.grad.clone()
-
-    energy_error = torch.abs((energy - energy_ref)/energy_ref)
-    grad_error = torch.max(torch.abs((grad - grad_ref)/grad_ref))
-
-    assert energy_error < 5e-7
-    assert grad_error < 5e-3
\ No newline at end of file
diff --git a/src/pytorch/TestOptimizedTorchANI.py b/src/pytorch/TestOptimizedTorchANI.py
deleted file mode 100644
index 5b1c2545..00000000
--- a/src/pytorch/TestOptimizedTorchANI.py
+++ /dev/null
@@ -1,138 +0,0 @@
-#
-# Copyright (c) 2020-2021 Acellera
-# Authors: Raimondas Galvelis
-#
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-#
-# The above copyright notice and this permission notice shall be included in all
-# copies or substantial portions of the Software.
-#
-# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-# SOFTWARE.
-#
-
-import mdtraj
-import os
-import pytest
-import tempfile
-import torch
-import torchani
-
-molecules = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'molecules')
-
-@pytest.mark.parametrize('deviceString', ['cpu', 'cuda'])
-@pytest.mark.parametrize('molFile', ['1hvj', '1hvk', '2iuz', '3hkw', '3hky', '3lka', '3o99'])
-def test_compare_with_native(deviceString, molFile):
-
-    if deviceString == 'cuda' and not torch.cuda.is_available():
-        pytest.skip('CUDA is not available')
-
-    from NNPOps import OptimizedTorchANI
-
-    device = torch.device(deviceString)
-
-    mol = mdtraj.load(os.path.join(molecules, f'{molFile}_ligand.mol2'))
-    atomicNumbers = torch.tensor([[atom.element.atomic_number for atom in mol.top.atoms]], device=device)
-    atomicPositions = torch.tensor(mol.xyz * 10, dtype=torch.float32, requires_grad=True, device=device)
-
-    nnp = torchani.models.ANI2x(periodic_table_index=True).to(device)
-    energy_ref = nnp((atomicNumbers, atomicPositions)).energies
-    energy_ref.backward()
-    grad_ref = atomicPositions.grad.clone()
-
-    nnp = OptimizedTorchANI(nnp, atomicNumbers).to(device)
-    energy = nnp((atomicNumbers, atomicPositions)).energies
-    atomicPositions.grad.zero_()
-    energy.backward()
-    grad = atomicPositions.grad.clone()
-
-    energy_error = torch.abs((energy - energy_ref)/energy_ref)
-    grad_error = torch.max(torch.abs((grad - grad_ref)/grad_ref))
-
-    assert energy_error < 5e-7
-    if molFile == '3o99':
-        assert grad_error < 7e-3
-    else:
-        assert grad_error < 5e-3
-
-@pytest.mark.parametrize('deviceString', ['cpu', 'cuda'])
-def test_compare_waterbox_pbc_with_native(deviceString):
-
-    if deviceString == 'cuda' and not torch.cuda.is_available():
-        pytest.skip('CUDA is not available')
-
-    from NNPOps import OptimizedTorchANI
-
-    device = torch.device(deviceString)
-
-    mol = mdtraj.load(os.path.join(molecules, 'water.pdb'))
-    atomicNumbers = torch.tensor([[atom.element.atomic_number for atom in mol.top.atoms]], device=device)
-    atomicPositions = torch.tensor(mol.xyz * 10, dtype=torch.float32, requires_grad=True, device=device)
-    cell = mol.unitcell_vectors[0]
-    cell = torch.tensor(cell, dtype=torch.float32, device=device)*10.0
-    pbc = torch.tensor([True, True, True], dtype=torch.bool, device=device)
-
-    nnp = torchani.models.ANI2x(periodic_table_index=True).to(device)
-    energy_ref = nnp((atomicNumbers, atomicPositions), cell=cell, pbc=pbc).energies
-    energy_ref.backward()
-    grad_ref = atomicPositions.grad.clone()
-
-    nnp = OptimizedTorchANI(nnp, atomicNumbers).to(device)
-    energy = nnp((atomicNumbers, atomicPositions), cell=cell, pbc=pbc).energies
-    atomicPositions.grad.zero_()
-    energy.backward()
-    grad = atomicPositions.grad.clone()
-
-    energy_error = torch.abs((energy - energy_ref)/energy_ref)
-    grad_error = torch.max(torch.abs((grad - grad_ref)/grad_ref))
-
-    assert energy_error < 5e-7
-    assert grad_error < 8e-3
-
-@pytest.mark.parametrize('deviceString', ['cpu', 'cuda'])
-@pytest.mark.parametrize('molFile', ['1hvj', '1hvk', '2iuz', '3hkw', '3hky', '3lka', '3o99'])
-def test_model_serialization(deviceString, molFile):
-
-    if deviceString == 'cuda' and not torch.cuda.is_available():
-        pytest.skip('CUDA is not available')
-
-    from NNPOps import OptimizedTorchANI
-
-    device = torch.device(deviceString)
-
-    mol = mdtraj.load(os.path.join(molecules, f'{molFile}_ligand.mol2'))
-    atomicNumbers = torch.tensor([[atom.element.atomic_number for atom in mol.top.atoms]], device=device)
-    atomicPositions = torch.tensor(mol.xyz * 10, dtype=torch.float32, requires_grad=True, device=device)
-
-    nnp_ref = torchani.models.ANI2x(periodic_table_index=True).to(device)
-    nnp_ref = OptimizedTorchANI(nnp_ref, atomicNumbers).to(device)
-
-    energy_ref = nnp_ref((atomicNumbers, atomicPositions)).energies
-    energy_ref.backward()
-    grad_ref = atomicPositions.grad.clone()
-
-    with tempfile.NamedTemporaryFile() as fd:
-
-        torch.jit.script(nnp_ref).save(fd.name)
-        nnp = torch.jit.load(fd.name)
-
-        energy = nnp((atomicNumbers, atomicPositions)).energies
-        atomicPositions.grad.zero_()
-        energy.backward()
-        grad = atomicPositions.grad.clone()
-
-    energy_error = torch.abs((energy - energy_ref)/energy_ref)
-    grad_error = torch.max(torch.abs((grad - grad_ref)/grad_ref))
-
-    assert energy_error < 5e-7
-    assert grad_error < 5e-3
\ No newline at end of file
diff --git a/src/pytorch/TestSpeciesConverter.py b/src/pytorch/TestSpeciesConverter.py
deleted file mode 100644
index da83ce22..00000000
--- a/src/pytorch/TestSpeciesConverter.py
+++ /dev/null
@@ -1,105 +0,0 @@
-#
-# Copyright (c) 2020-2021 Acellera
-# Authors: Raimondas Galvelis
-#
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-#
-# The above copyright notice and this permission notice shall be included in all
-# copies or substantial portions of the Software.
-#
-# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-# SOFTWARE.
-#
-
-import mdtraj
-import os
-import pytest
-import tempfile
-import torch
-import torchani
-
-molecules = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'molecules')
-
-def test_import():
-    import NNPOps
-    import NNPOps.SpeciesConverter
-
-@pytest.mark.parametrize('deviceString', ['cpu', 'cuda'])
-@pytest.mark.parametrize('molFile', ['1hvj', '1hvk', '2iuz', '3hkw', '3hky', '3lka', '3o99'])
-def test_compare_with_native(deviceString, molFile):
-
-    if deviceString == 'cuda' and not torch.cuda.is_available():
-        pytest.skip('CUDA is not available')
-
-    from NNPOps.SpeciesConverter import TorchANISpeciesConverter
-
-    device = torch.device(deviceString)
-
-    mol = mdtraj.load(os.path.join(molecules, f'{molFile}_ligand.mol2'))
-    atomicNumbers = torch.tensor([[atom.element.atomic_number for atom in mol.top.atoms]], device=device)
-    atomicPositions = torch.tensor(mol.xyz * 10, dtype=torch.float32, requires_grad=True, device=device)
-
-    nnp = torchani.models.ANI2x(periodic_table_index=True).to(device)
-    energy_ref = nnp((atomicNumbers, atomicPositions)).energies
-    energy_ref.backward()
-    grad_ref = atomicPositions.grad.clone()
-
-    nnp.species_converter = TorchANISpeciesConverter(nnp.species_converter, atomicNumbers).to(device)
-    energy = nnp((atomicNumbers, atomicPositions)).energies
-    atomicPositions.grad.zero_()
-    energy.backward()
-    grad = atomicPositions.grad.clone()
-
-    energy_error = torch.abs((energy - energy_ref)/energy_ref)
-    grad_error = torch.max(torch.abs((grad - grad_ref)/grad_ref))
-
-    assert energy_error < 5e-7
-    assert grad_error < 5e-3
-
-@pytest.mark.parametrize('deviceString', ['cpu', 'cuda'])
-@pytest.mark.parametrize('molFile', ['1hvj', '1hvk', '2iuz', '3hkw', '3hky', '3lka', '3o99'])
-def test_model_serialization(deviceString, molFile):
-
-    if deviceString == 'cuda' and not torch.cuda.is_available():
-        pytest.skip('CUDA is not available')
-
-    from NNPOps.SpeciesConverter import TorchANISpeciesConverter
-
-    device = torch.device(deviceString)
-
-    mol = mdtraj.load(os.path.join(molecules, f'{molFile}_ligand.mol2'))
-    atomicNumbers = torch.tensor([[atom.element.atomic_number for atom in mol.top.atoms]], device=device)
-    atomicPositions = torch.tensor(mol.xyz * 10, dtype=torch.float32, requires_grad=True, device=device)
-
-    nnp_ref = torchani.models.ANI2x(periodic_table_index=True).to(device)
-    nnp_ref.species_converter = TorchANISpeciesConverter(nnp_ref.species_converter, atomicNumbers).to(device)
-
-    energy_ref = nnp_ref((atomicNumbers, atomicPositions)).energies
-    energy_ref.backward()
-    grad_ref = atomicPositions.grad.clone()
-
-    with tempfile.NamedTemporaryFile() as fd:
-
-        torch.jit.script(nnp_ref).save(fd.name)
-        nnp = torch.jit.load(fd.name)
-
-        energy = nnp((atomicNumbers, atomicPositions)).energies
-        atomicPositions.grad.zero_()
-        energy.backward()
-        grad = atomicPositions.grad.clone()
-
-    energy_error = torch.abs((energy - energy_ref)/energy_ref)
-    grad_error = torch.max(torch.abs((grad - grad_ref)/grad_ref))
-
-    assert energy_error < 5e-7
-    assert grad_error < 5e-3
\ No newline at end of file
diff --git a/src/pytorch/TestSymmetryFunctions.py b/src/pytorch/TestSymmetryFunctions.py
index eae7c286..60364cd3 100644
--- a/src/pytorch/TestSymmetryFunctions.py
+++ b/src/pytorch/TestSymmetryFunctions.py
@@ -1,6 +1,6 @@
 #
-# Copyright (c) 2020 Acellera
-# Authors: Raimondas Galvelis
+# Copyright (c) 2020 Acellera, 2025 Stanford University and the Authors
+# Authors: Raimondas Galvelis, Evan Pretti
 #
 # Permission is hereby granted, free of charge, to any person obtaining a copy
 # of this software and associated documentation files (the "Software"), to deal
@@ -21,159 +21,103 @@
 # SOFTWARE.
 #
 
-import mdtraj
 import os
 import pytest
 import tempfile
 import torch
-import torchani
 
-molecules = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'molecules')
+test_data_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'test_data')
+
+VALUE_TOL = 1e-5
+GRADIENT_TOL = 1e-4
 
 def test_import():
     import NNPOps
     import NNPOps.SymmetryFunctions
 
 @pytest.mark.parametrize('deviceString', ['cpu', 'cuda'])
-@pytest.mark.parametrize('molFile', ['1hvj', '1hvk', '2iuz', '3hkw', '3hky', '3lka', '3o99'])
-def test_compare_with_native(deviceString, molFile):
-
-    if deviceString == 'cuda' and not torch.cuda.is_available():
-        pytest.skip('CUDA is not available')
-
-    from NNPOps.SymmetryFunctions import TorchANISymmetryFunctions
-
-    device = torch.device(deviceString)
-
-    mol = mdtraj.load(os.path.join(molecules, f'{molFile}_ligand.mol2'))
-    atomicNumbers = torch.tensor([[atom.element.atomic_number for atom in mol.top.atoms]], device=device)
-    atomicPositions = torch.tensor(mol.xyz * 10, dtype=torch.float32, requires_grad=True, device=device)
-
-    nnp = torchani.models.ANI2x(periodic_table_index=True).to(device)
-    energy_ref = nnp((atomicNumbers, atomicPositions)).energies
-    energy_ref.backward()
-    grad_ref = atomicPositions.grad.clone()
-
-    nnp.aev_computer = TorchANISymmetryFunctions(nnp.species_converter, nnp.aev_computer, atomicNumbers)
-    energy = nnp((atomicNumbers, atomicPositions)).energies
-    atomicPositions.grad.zero_()
-    energy.backward()
-    grad = atomicPositions.grad.clone()
-
-    energy_error = torch.abs((energy - energy_ref)/energy_ref)
-    grad_error = torch.max(torch.abs((grad - grad_ref)/grad_ref))
-
-    assert energy_error < 5e-7
-    if molFile == '3o99':
-        assert grad_error < 7e-3
-    else:
-        assert grad_error < 5e-3
-
-
-@pytest.mark.parametrize('deviceString', ['cpu', 'cuda'])
-def test_compare_waterbox_pbc_with_native(deviceString):
+@pytest.mark.parametrize('molFile', ['1hvj', '1hvk', '2iuz', '3hkw', '3hky', '3lka', '3o99', 'water'])
+def test_compare_with_reference(deviceString, molFile):
 
     if deviceString == 'cuda' and not torch.cuda.is_available():
         pytest.skip('CUDA is not available')
 
-    from NNPOps.SymmetryFunctions import TorchANISymmetryFunctions
+    from NNPOps.SymmetryFunctions import ANISymmetryFunctions
 
     device = torch.device(deviceString)
 
-    mol = mdtraj.load(os.path.join(molecules, 'water.pdb'))
-    atomicNumbers = torch.tensor([[atom.element.atomic_number for atom in mol.top.atoms]], device=device)
-    atomicPositions = torch.tensor(mol.xyz * 10, dtype=torch.float32, requires_grad=True, device=device)
-    cell = mol.unitcell_vectors[0]
-    cell = torch.tensor(cell, dtype=torch.float32, device=device)*10.0
-    pbc = torch.tensor([True, True, True], dtype=torch.bool, device=device)
-
-    nnp = torchani.models.ANI2x(periodic_table_index=True).to(device)
-    energy_ref = nnp((atomicNumbers, atomicPositions), cell=cell, pbc=pbc).energies
-    energy_ref.backward()
-    grad_ref = atomicPositions.grad.clone()
-
-    nnp.aev_computer = TorchANISymmetryFunctions(nnp.species_converter, nnp.aev_computer, atomicNumbers)
-    energy = nnp((atomicNumbers, atomicPositions), cell=cell, pbc=pbc).energies
-    atomicPositions.grad.zero_()
-    energy.backward()
-    grad = atomicPositions.grad.clone()
+    test_case = torch.load(os.path.join(test_data_path, f'{molFile}.pt'))
+    positions = test_case['positions'].to(device)
+    positions.requires_grad = True
+    cell = test_case['cell']
+    if cell is not None:
+        cell = cell.to(device)
 
-    energy_error = torch.abs((energy - energy_ref)/energy_ref)
-    grad_error = torch.max(torch.abs((grad - grad_ref)/grad_ref))
+    expected = test_case['output'].to(device)
+    actual = torch.concat(ANISymmetryFunctions(**test_case['parameters'])(positions, cell), dim=1)
+    total = torch.sum(actual)
+    total.backward()
 
-    assert energy_error < 5e-7
-    assert grad_error < 7.5e-3
+    assert torch.allclose(actual, expected, rtol=VALUE_TOL, atol=VALUE_TOL)
+    assert torch.allclose(positions.grad, test_case['grad'].to(device), rtol=GRADIENT_TOL, atol=GRADIENT_TOL)
 
 @pytest.mark.parametrize('deviceString', ['cpu', 'cuda'])
-@pytest.mark.parametrize('molFile', ['1hvj', '1hvk', '2iuz', '3hkw', '3hky', '3lka', '3o99'])
+@pytest.mark.parametrize('molFile', ['1hvj', '1hvk', '2iuz', '3hkw', '3hky', '3lka', '3o99', 'water'])
 def test_model_serialization(deviceString, molFile):
 
     if deviceString == 'cuda' and not torch.cuda.is_available():
         pytest.skip('CUDA is not available')
 
-    from NNPOps.SymmetryFunctions import TorchANISymmetryFunctions
+    from NNPOps.SymmetryFunctions import ANISymmetryFunctions
 
     device = torch.device(deviceString)
 
-    mol = mdtraj.load(os.path.join(molecules, f'{molFile}_ligand.mol2'))
-    atomicNumbers = torch.tensor([[atom.element.atomic_number for atom in mol.top.atoms]], device=device)
-    atomicPositions = torch.tensor(mol.xyz * 10, dtype=torch.float32, requires_grad=True, device=device)
+    test_case = torch.load(os.path.join(test_data_path, f'{molFile}.pt'))
+    positions = test_case['positions'].to(device)
+    positions.requires_grad = True
+    cell = test_case['cell']
+    if cell is not None:
+        cell = cell.to(device)
 
-    nnp_ref = torchani.models.ANI2x(periodic_table_index=True).to(device)
-    nnp_ref.aev_computer = TorchANISymmetryFunctions(nnp_ref.species_converter, nnp_ref.aev_computer, atomicNumbers)
-
-    energy_ref = nnp_ref((atomicNumbers, atomicPositions)).energies
-    energy_ref.backward()
-    grad_ref = atomicPositions.grad.clone()
+    expected = test_case['output'].to(device)
 
     with tempfile.NamedTemporaryFile() as fd:
 
-        torch.jit.script(nnp_ref).save(fd.name)
-        nnp = torch.jit.load(fd.name)
-
-        energy = nnp((atomicNumbers, atomicPositions)).energies
-        atomicPositions.grad.zero_()
-        energy.backward()
-        grad = atomicPositions.grad.clone()
+        torch.jit.script(ANISymmetryFunctions(**test_case['parameters'])).save(fd.name)
+        actual = torch.concat(torch.jit.load(fd.name)(positions, cell), dim=1)
 
-    energy_error = torch.abs((energy - energy_ref)/energy_ref)
-    grad_error = torch.max(torch.abs((grad - grad_ref)/grad_ref))
+    total = torch.sum(actual)
+    total.backward()
 
-    assert energy_error < 5e-7
-    assert grad_error < 5e-3
+    assert torch.allclose(actual, expected, rtol=VALUE_TOL, atol=VALUE_TOL)
+    assert torch.allclose(positions.grad, test_case['grad'].to(device), rtol=GRADIENT_TOL, atol=GRADIENT_TOL)
 
-@pytest.mark.parametrize('molFile', ['1hvj', '1hvk', '2iuz', '3hkw', '3hky', '3lka', '3o99'])
+@pytest.mark.parametrize('molFile', ['1hvj', '1hvk', '2iuz', '3hkw', '3hky', '3lka', '3o99', 'water'])
 def test_non_default_stream(molFile):
 
     if not torch.cuda.is_available():
         pytest.skip('CUDA is not available')
 
-    from NNPOps.SymmetryFunctions import TorchANISymmetryFunctions
+    from NNPOps.SymmetryFunctions import ANISymmetryFunctions
 
     device = torch.device('cuda')
 
-    mol = mdtraj.load(os.path.join(molecules, f'{molFile}_ligand.mol2'))
-    atomicNumbers = torch.tensor([[atom.element.atomic_number for atom in mol.top.atoms]], device=device)
-    atomicPositions = torch.tensor(mol.xyz * 10, dtype=torch.float32, requires_grad=True, device=device)
-
-    nnp = torchani.models.ANI2x(periodic_table_index=True).to(device)
-    nnp.aev_computer = TorchANISymmetryFunctions(nnp.species_converter, nnp.aev_computer, atomicNumbers)
-
-    energy_ref = nnp((atomicNumbers, atomicPositions)).energies
-    energy_ref.backward()
-    grad_ref = atomicPositions.grad.clone()
+    test_case = torch.load(os.path.join(test_data_path, f'{molFile}.pt'))
+    positions = test_case['positions'].to(device)
+    positions.requires_grad = True
+    cell = test_case['cell']
+    if cell is not None:
+        cell = cell.to(device)
 
+    expected = test_case['output'].to(device)
+    module = ANISymmetryFunctions(**test_case['parameters'])
     stream = torch.cuda.Stream()
     stream.wait_stream(torch.cuda.current_stream())
     with torch.cuda.stream(stream):
-        energy = nnp((atomicNumbers, atomicPositions)).energies
-        atomicPositions.grad.zero_()
-        energy.backward()
-        grad = atomicPositions.grad.clone()
+        actual = torch.concat(module(positions, cell), dim=1)
+        total = torch.sum(actual)
+        total.backward()
     torch.cuda.current_stream().wait_stream(stream)
 
-    energy_error = torch.abs((energy - energy_ref)/energy_ref)
-    grad_error = torch.max(torch.abs((grad - grad_ref)/grad_ref))
-
-    assert energy_error < 5e-7
-    assert grad_error < 5e-3
\ No newline at end of file
+    assert torch.allclose(actual, expected, rtol=5e-5)
+    assert torch.allclose(positions.grad, test_case['grad'].to(device), rtol=GRADIENT_TOL, atol=GRADIENT_TOL)
diff --git a/src/pytorch/__init__.py b/src/pytorch/__init__.py
index 5e06bd5b..430fec2e 100644
--- a/src/pytorch/__init__.py
+++ b/src/pytorch/__init__.py
@@ -12,6 +12,3 @@
     pass
 
 torch.ops.load_library(os.path.join(os.path.dirname(__file__), 'libNNPOpsPyTorch.so'))
-
-
-from NNPOps.OptimizedTorchANI import OptimizedTorchANI
diff --git a/src/pytorch/neighbors/getNeighborPairs.py b/src/pytorch/neighbors/getNeighborPairs.py
index 12a4b03c..e84341a5 100644
--- a/src/pytorch/neighbors/getNeighborPairs.py
+++ b/src/pytorch/neighbors/getNeighborPairs.py
@@ -141,7 +141,8 @@ def getNeighborPairs(
 
     if box_vectors is None:
         box_vectors = empty((0, 0), device=positions.device, dtype=positions.dtype)
+    # Convert max_num_pairs to a Python int in case it is a NumPy integer.
     neighbors, deltas, distances, number_found_pairs = ops.neighbors.getNeighborPairs(
-        positions, cutoff, max_num_pairs, box_vectors, check_errors
+        positions, cutoff, int(max_num_pairs), box_vectors, check_errors
     )
     return neighbors, deltas, distances, number_found_pairs
diff --git a/src/pytorch/test_data/1hvj.pt b/src/pytorch/test_data/1hvj.pt
new file mode 100644
index 00000000..6e8ff2df
Binary files /dev/null and b/src/pytorch/test_data/1hvj.pt differ
diff --git a/src/pytorch/test_data/1hvk.pt b/src/pytorch/test_data/1hvk.pt
new file mode 100644
index 00000000..9ea3a662
Binary files /dev/null and b/src/pytorch/test_data/1hvk.pt differ
diff --git a/src/pytorch/test_data/2iuz.pt b/src/pytorch/test_data/2iuz.pt
new file mode 100644
index 00000000..0d0b93af
Binary files /dev/null and b/src/pytorch/test_data/2iuz.pt differ
diff --git a/src/pytorch/test_data/3hkw.pt b/src/pytorch/test_data/3hkw.pt
new file mode 100644
index 00000000..8c1381bd
Binary files /dev/null and b/src/pytorch/test_data/3hkw.pt differ
diff --git a/src/pytorch/test_data/3hky.pt b/src/pytorch/test_data/3hky.pt
new file mode 100644
index 00000000..91a7cf90
Binary files /dev/null and b/src/pytorch/test_data/3hky.pt differ
diff --git a/src/pytorch/test_data/3lka.pt b/src/pytorch/test_data/3lka.pt
new file mode 100644
index 00000000..3822f124
Binary files /dev/null and b/src/pytorch/test_data/3lka.pt differ
diff --git a/src/pytorch/test_data/3o99.pt b/src/pytorch/test_data/3o99.pt
new file mode 100644
index 00000000..c7d0972b
Binary files /dev/null and b/src/pytorch/test_data/3o99.pt differ
diff --git a/src/pytorch/test_data/water.pt b/src/pytorch/test_data/water.pt
new file mode 100644
index 00000000..2838e5d8
Binary files /dev/null and b/src/pytorch/test_data/water.pt differ
diff --git a/src/schnet/CudaCFConv.cu b/src/schnet/CudaCFConv.cu
index 3bf0e28c..45b62992 100644
--- a/src/schnet/CudaCFConv.cu
+++ b/src/schnet/CudaCFConv.cu
@@ -134,7 +134,7 @@ void CudaCFConvNeighbors::build(const float* positions, const float* periodicBox
 
     cudaPointerAttributes attrib;
     cudaError_t result = cudaPointerGetAttributes(&attrib, positions);
-    if (result != cudaSuccess || attrib.devicePointer == 0) {
+    if (result != cudaSuccess || attrib.type != cudaMemoryTypeDevice || attrib.devicePointer == 0) {
         CHECK_RESULT(cudaMemcpyAsync(this->positions, positions, 3*getNumAtoms()*sizeof(float), cudaMemcpyDefault));
         devicePositions = this->positions;
     }
@@ -147,7 +147,7 @@ void CudaCFConvNeighbors::build(const float* positions, const float* periodicBox
    const float* hostBoxVectors;
     if (getPeriodic()) {
         result = cudaPointerGetAttributes(&attrib, periodicBoxVectors);
-        if (result != cudaSuccess || attrib.devicePointer == 0) {
+        if (result != cudaSuccess || attrib.type != cudaMemoryTypeDevice || attrib.devicePointer == 0) {
             CHECK_RESULT(cudaMemcpyAsync(this->periodicBoxVectors, periodicBoxVectors, 9*sizeof(float), cudaMemcpyDefault));
             hostBoxVectors = periodicBoxVectors;
             deviceBoxVectors = this->periodicBoxVectors;
@@ -239,7 +239,7 @@ CudaCFConv::~CudaCFConv() {
 float* CudaCFConv::ensureOnDevice(float* arg, float*& deviceMemory, int size) {
     cudaPointerAttributes attrib;
     cudaError_t result = cudaPointerGetAttributes(&attrib, arg);
-    if (result != cudaSuccess || attrib.devicePointer == 0) {
+    if (result != cudaSuccess || attrib.type != cudaMemoryTypeDevice || attrib.devicePointer == 0) {
         if (deviceMemory == 0)
             CHECK_RESULT(cudaMallocManaged(&deviceMemory, size));
         CHECK_RESULT(cudaMemcpyAsync(deviceMemory, arg, size, cudaMemcpyDefault));
@@ -251,7 +251,7 @@ float* CudaCFConv::ensureOnDevice(float* arg, float*& deviceMemory, int size) {
 const float* CudaCFConv::ensureOnDevice(const float* arg, float*& deviceMemory, int size) {
     cudaPointerAttributes attrib;
     cudaError_t result = cudaPointerGetAttributes(&attrib, arg);
-    if (result != cudaSuccess || attrib.devicePointer == 0) {
+    if (result != cudaSuccess || attrib.type != cudaMemoryTypeDevice || attrib.devicePointer == 0) {
         if (deviceMemory == 0)
             CHECK_RESULT(cudaMallocManaged(&deviceMemory, size));
         CHECK_RESULT(cudaMemcpyAsync(deviceMemory, arg, size, cudaMemcpyDefault));
diff --git a/src/setup.py b/src/setup.py
index f75eac43..b3c32707 100644
--- a/src/setup.py
+++ b/src/setup.py
@@ -35,7 +35,6 @@ def set_torch_cuda_arch_list():
 
 sources = [
     opj("ani", "CpuANISymmetryFunctions.cpp"),
-    opj("pytorch", "BatchedNN.cpp"),
     opj("pytorch", "CFConv.cpp"),
     opj("pytorch", "CFConvNeighbors.cpp"),
     opj("pytorch", "SymmetryFunctions.cpp"),