Add missing cached_weights line

30e051cf · Gustaf Ahdritz · GitHub · 14cce39a · 30e051cf
Unverified Commit 30e051cf authored Dec 31, 2021 by Gustaf Ahdritz Committed by GitHub Dec 31, 2021
Show whitespace changes
Inline Side-by-side

Showing with 2 additions and 0 deletions

train_openfold.py train_openfold.py +2 -0

No files found.
--- a/train_openfold.py
+++ b/train_openfold.py
@@ -49,6 +49,8 @@ class OpenFoldWrapper(pl.LightningModule):
            model=self.model, decay=config.ema.decay
        )
        
+        self.cached_weights = None
+
    def forward(self, batch):
        return self.model(batch)