Add skip_loss for decode-only mode: no model build, no loss, no checkpoint

ConnectomicsModule(cfg, model=Identity(), skip_loss=True) skips all
loss/optimizer/weighter construction. main.py uses this when
inference.saved_prediction_path is set.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: Donglai Wei

connectomics/training/lightning/model.py

@@ -90,6 +90,7 @@ def __init__(
         self,
         cfg: Union[Config, DictConfig],
         model: Optional[nn.Module] = None,
+        skip_loss: bool = False,
     ):
         super().__init__()
         self.cfg = cfg
@@ -98,6 +99,17 @@ def __init__(
         # Build model
         self.model = model if model is not None else self._build_model(cfg)
 
+        # Skip loss/optimizer setup for decode-only mode
+        if skip_loss:
+            self.loss_functions = []
+            self.loss_weights = []
+            self.loss_metadata = []
+            self.loss_weighter = None
+            self.enable_nan_detection = False
+            self.debug_on_nan = False
+            self.loss_orchestrator = None
+            return
+
         # Build loss functions
         self.loss_functions = self._build_losses(cfg)
         self.loss_weights = self._extract_loss_weights(cfg)

scripts/main.py

@@ -904,7 +904,13 @@ def main():
     )
 
     # Create model
-    if tta_cached:
+    if has_saved_prediction:
+        print(f"  Decode-only mode: loading predictions from {_saved_pred}")
+        print(f"  Skipping model build entirely.")
+        model = ConnectomicsModule(cfg, model=torch.nn.Identity(), skip_loss=True)
+        model._skip_inference = True
+        ckpt_path = None
+    elif tta_cached:
         print(
             f"  Cached intermediate predictions found; "
             f"creating lightweight module (skipping {cfg.model.arch.type} build)."