fix doc build and test failure

Author: Donglai Wei

connectomics/config/auto_config.py

@@ -390,7 +390,9 @@ def auto_plan_config(
         if getattr(config.optimization, "precision", None) is not None:
             manual_overrides["precision"] = config.optimization.precision
         if getattr(config.optimization, "accumulate_grad_batches", None) is not None:
-            manual_overrides["accumulate_grad_batches"] = config.optimization.accumulate_grad_batches
+            manual_overrides["accumulate_grad_batches"] = (
+                config.optimization.accumulate_grad_batches
+            )
 
         opt_cfg = getattr(config.optimization, "optimizer", None)
         if opt_cfg and getattr(opt_cfg, "lr", None) is not None:
@@ -406,8 +408,7 @@ def auto_plan_config(
 
     # Plan
     use_mixed_precision = not (
-        hasattr(config, "optimization")
-        and getattr(config.optimization, "precision", None) == "32"
+        hasattr(config, "optimization") and getattr(config.optimization, "precision", None) == "32"
     )
 
     result = planner.plan(

connectomics/config/hydra_config.py

@@ -1220,6 +1220,18 @@ class TuneConfig:
                 DecodingParameterSpace,
                 PostprocessingParameterSpace,
                 ParameterSpaceConfig,
+                # Core config dataclasses (for Lightning checkpoints)
+                Config,
+                SystemConfig,
+                SystemTrainingConfig,
+                SystemInferenceConfig,
+                ModelConfig,
+                DataConfig,
+                OptimizationConfig,
+                MonitorConfig,
+                InferenceConfig,
+                TestConfig,
+                TuneConfig,
             ]
         )
 except Exception:

connectomics/data/augment/monai_transforms.py

@@ -237,7 +237,9 @@ def _apply_missing_section(
         if is_tensor:
             keep_mask = torch.ones(depth, dtype=torch.bool, device=img.device)
             keep_mask[indices_to_remove] = False
-            return torch.index_select(img, dim=depth_axis, index=keep_mask.nonzero(as_tuple=False).squeeze(-1))
+            return torch.index_select(
+                img, dim=depth_axis, index=keep_mask.nonzero(as_tuple=False).squeeze(-1)
+            )
         else:
             return np.delete(img, indices_to_remove, axis=depth_axis)
 

connectomics/training/lit/model.py

@@ -82,12 +82,10 @@ def __init__(
             if hasattr(cfg.model, "loss_weights")
             else [1.0] * len(self.loss_functions)
         )
-        self.multi_task_config = (
-            getattr(cfg.model, "multi_task_config", None) or []
-        )
+        self.multi_task_config = getattr(cfg.model, "multi_task_config", None) or []
         self.multi_task_enabled = len(self.multi_task_config) > 0
-        num_tasks = len(self.multi_task_config) if self.multi_task_config else len(
-            self.loss_functions
+        num_tasks = (
+            len(self.multi_task_config) if self.multi_task_config else len(self.loss_functions)
         )
         self.loss_weighter = build_loss_weighter(cfg, num_tasks=num_tasks, model=self.model)
 

connectomics/training/lit/trainer.py

@@ -24,12 +24,38 @@
 from pytorch_lightning.strategies import DDPStrategy
 
 from ...config import Config
+from ...config.hydra_config import (
+    SystemConfig,
+    SystemTrainingConfig,
+    SystemInferenceConfig,
+    ModelConfig,
+    DataConfig,
+    OptimizationConfig,
+    MonitorConfig,
+    InferenceConfig,
+    TestConfig,
+    TuneConfig,
+)
 from .callbacks import VisualizationCallback, EMAWeightsCallback
 
 # Register safe globals for PyTorch 2.6+ checkpoint loading
 # This allows our Config class to be unpickled from Lightning checkpoints
 try:
-    torch.serialization.add_safe_globals([Config])
+    torch.serialization.add_safe_globals(
+        [
+            Config,
+            SystemConfig,
+            SystemTrainingConfig,
+            SystemInferenceConfig,
+            ModelConfig,
+            DataConfig,
+            OptimizationConfig,
+            MonitorConfig,
+            InferenceConfig,
+            TestConfig,
+            TuneConfig,
+        ]
+    )
 except AttributeError:
     # PyTorch < 2.6 doesn't have add_safe_globals
     pass

docs/source/_templates/layout.html

@@ -45,7 +45,8 @@
 
   {% endif %}
 
-  <link rel="stylesheet" href="{{ pathto('_static/' + style, 1) }}" type="text/css" />
+  {# `style` is absent on newer Sphinx, so default to the theme stylesheet when undefined #}
+  <link rel="stylesheet" href="{{ pathto('_static/' + (style | default('css/theme.css')), 1) }}" type="text/css" />
   <!-- <link rel="stylesheet" href="{{ pathto('_static/pygments.css', 1) }}" type="text/css" /> -->
   {%- for css in css_files %}
   {%- if css|attr("rel") %}