diff --git a/src/metatrain/experimental/pet/tests/test_pet_compatibility.py b/src/metatrain/experimental/pet/tests/test_pet_compatibility.py
index 16f7955a..852750be 100644
--- a/src/metatrain/experimental/pet/tests/test_pet_compatibility.py
+++ b/src/metatrain/experimental/pet/tests/test_pet_compatibility.py
@@ -159,7 +159,7 @@ def test_predictions_compatibility(cutoff):
         "neighbors_pos": batch.neighbors_pos,
     }
 
-    pet = model._module.pet
+    pet = model.module.pet
 
     pet_prediction = pet.forward(batch_dict)
 
diff --git a/src/metatrain/utils/llpr.py b/src/metatrain/utils/llpr.py
index 20eb68df..7f23fda9 100644
--- a/src/metatrain/utils/llpr.py
+++ b/src/metatrain/utils/llpr.py
@@ -31,7 +31,7 @@ def __init__(
         super().__init__()
 
         self.model = model
-        self.ll_feat_size = self.model._module.last_layer_feature_size
+        self.ll_feat_size = self.model.module.last_layer_feature_size
 
         # update capabilities: now we have additional outputs for the uncertainty
         old_capabilities = self.model.capabilities()