Fix bug in training script

41c68b1e · Gustaf Ahdritz · 1d657f55 · 41c68b1e
Commit 41c68b1e authored Nov 15, 2021 by Gustaf Ahdritz
Show whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

train_openfold.py train_openfold.py +1 -1

No files found.
--- a/train_openfold.py
+++ b/train_openfold.py
@@ -73,7 +73,7 @@ class OpenFoldWrapper(pl.LightningModule):
    def validation_step(self, batch, batch_idx):
        # At the start of validation, load the EMA weights
        if(self.cached_weights is None):
-            self.cached_weights = model.state_dict()
+            self.cached_weights = self.model.state_dict()
            self.model.load_state_dict(self.ema.state_dict()["params"])
        
        # Calculate validation loss