fixed test_models.py and tensornet parameters

choderalab · Sep 10, 2024 · 2d713b0 · 2d713b0
1 parent 47c4d3b
commit 2d713b0
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/modelforge/potential/models.py b/modelforge/potential/models.py
@@ -558,7 +558,7 @@ def __init__(self, box_vectors: torch.Tensor, periodic: bool):
             [box_vectors[0][0], box_vectors[1][1], box_vectors[2][2]]
         ).to(box_vectors.device)
 
-        self.register_buffer("cutoff", torch.tensor(periodic))
+        self.periodic = periodic
 
     def forward(self, coordinate_i: torch.Tensor, coordinate_j: torch.Tensor):
         r_ij = coordinate_i - coordinate_j

diff --git a/modelforge/potential/parameters.py b/modelforge/potential/parameters.py
@@ -199,7 +199,7 @@ class CoreParameter(ParametersBase):
         number_of_radial_basis_functions: int
         maximum_interaction_radius: float
         minimum_interaction_radius: float
-        highest_atomic_number: int
+        maximum_atomic_number: int
         equivariance_invariance_group: str
         activation_function_parameter: ActivationFunctionConfig
 

diff --git a/modelforge/potential/tensornet.py b/modelforge/potential/tensornet.py
@@ -170,7 +170,7 @@ def __init__(
         number_of_radial_basis_functions: int,
         maximum_interaction_radius: float,
         minimum_interaction_radius: float,
-        highest_atomic_number: int,
+        maximum_atomic_number: int,
         equivariance_invariance_group: str,
         activation_function_parameter: Dict,
         potential_seed: int = -1,