Properly save/load criterion

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

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"
 

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

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"
 

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"
 

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":