diff --git a/modelforge/potential/potential.py b/modelforge/potential/potential.py
index 2af6f769..fafa6c09 100644
--- 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
index 9d9f444f..face5172 100644
--- 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"],