Merge pull request #101 from a-r-j/amorehead-decoder-disable

Add ability for `BenchmarkModel` to have its decoder disabled

Author: amorehead

CHANGELOG.md

@@ -15,6 +15,7 @@
 * Add support for handling training/validation OOMs gracefully [#81](https://github.com/a-r-j/ProteinWorkshop/pull/81)
 * Add support for handling backward OOMs gracefully [#83](https://github.com/a-r-j/ProteinWorkshop/pull/83)
 * Update GCPNet paper link [#85](https://github.com/a-r-j/ProteinWorkshop/pull/85)
+* Add ability for `BenchmarkModel` to have its decoder disabled [#101](https://github.com/a-r-j/ProteinWorkshop/pull/101)
 
 ### Framework
 

proteinworkshop/models/base.py

@@ -412,9 +412,13 @@ def __init__(self, cfg: DictConfig) -> None:
         self.encoder: nn.Module = hydra.utils.instantiate(cfg.encoder)
         logger.info(self.encoder)
 
-        logger.info("Instantiating decoders...")
-        self.decoder: nn.ModuleDict = self._build_output_decoders()
-        logger.info(self.decoder)
+        if hasattr(cfg.decoder, "disable") and cfg.decoder.disable:
+            logger.info("Disabling decoder as requested")
+            self.decoder = None
+        else:
+            logger.info("Instantiating decoders...")
+            self.decoder: nn.ModuleDict = self._build_output_decoders()
+            logger.info(self.decoder)
 
         logger.info("Instantiating losses...")
         self.losses = self.configure_losses(cfg.task.losses)