From 8ab62eafff9fc4c3357fff0d9767c51982c2cec6 Mon Sep 17 00:00:00 2001
From: lukasugar <lukasugar@gmail.com>
Date: Wed, 13 May 2026 14:15:16 +0200
Subject: [PATCH] Expose disabling decoder laoding for cls and reg tasks

---
 asparagus/modules/lightning_modules/clsreg_module.py | 2 ++
 asparagus/pipeline/run/finetune_cls.py               | 1 +
 asparagus/pipeline/run/finetune_reg.py               | 1 +
 configs/default_finetune_cls.yaml                    | 3 ++-
 configs/default_finetune_reg.yaml                    | 3 ++-
 5 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/asparagus/modules/lightning_modules/clsreg_module.py b/asparagus/modules/lightning_modules/clsreg_module.py
index 6d4838c..1d04326 100644
--- a/asparagus/modules/lightning_modules/clsreg_module.py
+++ b/asparagus/modules/lightning_modules/clsreg_module.py
@@ -33,6 +33,7 @@ def __init__(
         momentum: float = 0.99,
         log_image_every_n_epochs: int = 50,
         test_output_path: str = None,
+        load_decoder: bool = True,
         repeat_stem_weights: bool = True,
     ):
         super().__init__(
@@ -50,6 +51,7 @@ def __init__(
             weight_decay=weight_decay,
             nesterov=nesterov,
             momentum=momentum,
+            load_decoder=load_decoder,
             repeat_stem_weights=repeat_stem_weights,
         )
         self.loss = None
diff --git a/asparagus/pipeline/run/finetune_cls.py b/asparagus/pipeline/run/finetune_cls.py
index 4c2569e..0e9bd0e 100644
--- a/asparagus/pipeline/run/finetune_cls.py
+++ b/asparagus/pipeline/run/finetune_cls.py
@@ -114,6 +114,7 @@ def main(cfg: DictConfig) -> None:
         log_image_every_n_epochs=cfg.logger.log_images_every_n_epoch,
         optimizer=cfg.model.finetune_optim,
         learning_rate=cfg.model.finetune_lr,
+        load_decoder=cfg.training.load_decoder,
         repeat_stem_weights=cfg.training.repeat_stem_weights,
         test_output_path=os.path.join(
             path_store.run_dir,
diff --git a/asparagus/pipeline/run/finetune_reg.py b/asparagus/pipeline/run/finetune_reg.py
index 4388726..8c37f4c 100644
--- a/asparagus/pipeline/run/finetune_reg.py
+++ b/asparagus/pipeline/run/finetune_reg.py
@@ -114,6 +114,7 @@ def main(cfg: DictConfig) -> None:
         learning_rate=cfg.model.finetune_lr,
         warmup_epochs=cfg.training.warmup_epochs,
         weight_decay=cfg.model.weight_decay,
+        load_decoder=cfg.training.load_decoder,
         test_output_path=os.path.join(
             path_store.run_dir,
             "predictions",
diff --git a/configs/default_finetune_cls.yaml b/configs/default_finetune_cls.yaml
index 3514005..8d2b0ec 100644
--- a/configs/default_finetune_cls.yaml
+++ b/configs/default_finetune_cls.yaml
@@ -35,6 +35,7 @@ training:
   mask_ratio: 0.0
   check_val_every_n_epoch: 3
   accumulate_grad_batches: 1
+  load_decoder: True
   repeat_stem_weights: True
 
 transforms:
@@ -53,4 +54,4 @@ logger:
   mlflow_logging: False
   log_every_n_steps: 250
   log_images_every_n_epoch: 5
-  log_to_stdout: True
\ No newline at end of file
+  log_to_stdout: True
diff --git a/configs/default_finetune_reg.yaml b/configs/default_finetune_reg.yaml
index ca653d9..572d755 100644
--- a/configs/default_finetune_reg.yaml
+++ b/configs/default_finetune_reg.yaml
@@ -32,6 +32,7 @@ training:
   accumulate_grad_batches: 1
   warmup_epochs: 0
   check_val_every_n_epoch: 1
+  load_decoder: True
 
 transforms:
   cpu_tr_transforms: CPU_clsreg_train_transforms_crop
@@ -49,4 +50,4 @@ logger:
   mlflow_logging: False
   log_every_n_steps: 250
   log_images_every_n_epoch: 5
-  log_to_stdout: True
\ No newline at end of file
+  log_to_stdout: True