remove remaining keyword use

choderalab · Oct 22, 2024 · 69b3dd4 · 69b3dd4
1 parent 9e02af5
commit 69b3dd4
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 3 deletions.
diff --git a/modelforge/potential/potential.py b/modelforge/potential/potential.py
@@ -782,7 +782,6 @@ def load_inference_model_from_checkpoint(checkpoint_path: str) -> Potential:
 
     # Create the model in inference mode
     model = NeuralNetworkPotentialFactory.generate_potential(
-        use="inference",
         potential_parameter=potential_parameter,
         dataset_statistic=dataset_statistic,
         potential_seed=potential_seed,

diff --git a/modelforge/tests/test_schnet.py b/modelforge/tests/test_schnet.py
@@ -20,8 +20,7 @@ def setup_schnet_model(potential_seed: Optional[int] = None):
     config["potential"].core_parameter.number_of_radial_basis_functions = 5
     config["potential"].core_parameter.number_of_filters = 12
 
-    model = NeuralNetworkPotentialFactory.generate_potential(
-        use="training",
+    model = NeuralNetworkPotentialFactory.generate_trainer(
         potential_parameter=config["potential"],
         training_parameter=config["training"],
         dataset_parameter=config["dataset"],