diff --git a/src/ptbench/models/alexnet.py b/src/ptbench/models/alexnet.py
index 7aaaccb6a18e12e5e0d03c00e9f0974d470b31df..59acba158565723ab28a483e53b87b1395de742f 100644
--- a/src/ptbench/models/alexnet.py
+++ b/src/ptbench/models/alexnet.py
@@ -26,7 +26,7 @@ class Alexnet(pl.LightningModule):
     ):
         super().__init__()
 
-        self.save_hyperparameters()
+        self.save_hyperparameters(ignore=["criterion", "criterion_valid"])
 
         self.criterion = criterion
         self.criterion_valid = criterion_valid
diff --git a/src/ptbench/models/densenet.py b/src/ptbench/models/densenet.py
index 4e5b34c09331ed0140eb49afaf812a4c84e59b56..b44dac93f46447ef9fa8f3fbf0ca8c9f13163812 100644
--- a/src/ptbench/models/densenet.py
+++ b/src/ptbench/models/densenet.py
@@ -27,7 +27,7 @@ class Densenet(pl.LightningModule):
     ):
         super().__init__()
 
-        self.save_hyperparameters()
+        self.save_hyperparameters(ignore=["criterion", "criterion_valid"])
 
         self.name = "Densenet"
 
diff --git a/src/ptbench/models/logistic_regression.py b/src/ptbench/models/logistic_regression.py
index 684155b41e2e92a4903ae7ffbbe67dce74d93d70..ad56cb80530b3721e7aae20f6d3ebf03e6c19250 100644
--- a/src/ptbench/models/logistic_regression.py
+++ b/src/ptbench/models/logistic_regression.py
@@ -20,7 +20,7 @@ class LogisticRegression(pl.LightningModule):
     ):
         super().__init__()
 
-        self.save_hyperparameters()
+        self.save_hyperparameters(ignore=["criterion", "criterion_valid"])
 
         self.criterion = criterion
         self.criterion_valid = criterion_valid
diff --git a/src/ptbench/models/pasa.py b/src/ptbench/models/pasa.py
index 8c9705e61c474b0f029dd0418a067e4ded593820..af47d9e3d96afecde176ea193c9e0d449f341ee5 100644
--- a/src/ptbench/models/pasa.py
+++ b/src/ptbench/models/pasa.py
@@ -38,7 +38,7 @@ class PASA(pl.LightningModule):
     ):
         super().__init__()
 
-        self.save_hyperparameters()
+        self.save_hyperparameters(ignore=["criterion", "criterion_valid"])
 
         self.name = "pasa"
 
diff --git a/src/ptbench/models/signs_to_tb.py b/src/ptbench/models/signs_to_tb.py
index 653b590a380327e99e860dba2e10e8c72d17282d..0169a1b8fa008786829a1f301260efe3d695df7e 100644
--- a/src/ptbench/models/signs_to_tb.py
+++ b/src/ptbench/models/signs_to_tb.py
@@ -20,7 +20,7 @@ class SignsToTB(pl.LightningModule):
     ):
         super().__init__()
 
-        self.save_hyperparameters()
+        self.save_hyperparameters(ignore=["criterion", "criterion_valid"])
 
         self.name = "signs_to_tb"
 
diff --git a/src/ptbench/scripts/predict.py b/src/ptbench/scripts/predict.py
index 82939f25fa186bff1e6fea0636a800cdcbba2ac6..860d95b293f22895e62ebc825a03550d24018806 100644
--- a/src/ptbench/scripts/predict.py
+++ b/src/ptbench/scripts/predict.py
@@ -122,7 +122,9 @@ def predict(
 
     dataset = dataset if isinstance(dataset, dict) else dict(test=dataset)
 
-    model = model.load_from_checkpoint(weight)
+    model = model.load_from_checkpoint(
+        weight, criterion=model.criterion, criterion_valid=model.criterion_valid
+    )
 
     # Logistic regressor weights
     if model.name == "logistic_regression":