[main] Integration/peerlab

.gitignore

@@ -206,4 +206,5 @@ __marimo__/
 .dev
 .dev/*
 *.pyc
-*memray*
+*memray*
+nvidia/

src/segger/__init__.py

@@ -13,11 +13,6 @@
 torch.cuda.memory.change_current_allocator(rmm_torch_allocator)
 enable_statistics()
 
-# Apply pytorch patches for issue pytorch/pytorch#51871 (CUDA nonzero INT_MAX limit).
-# Must run BEFORE any segger module imports HeteroData / bipartite_subgraph.
-from ._patches import apply as _apply_patches
-_apply_patches()
-
 def free_mem_str() -> str:
     stats = get_statistics()
     return (

src/segger/cli/debug.py

@@ -9,6 +9,7 @@
 
 from ..debug.segmentation import run_segmentation_only
 from ..debug.prediction import run_prediction_only
+from ..utils import setup_logging
 
 debug = App(name="debug", help="Utilities for debugging and testing individual components.")
 
@@ -41,6 +42,7 @@ def predict_only_cli(
     )],
 ):
     """Run prediction only."""
+    setup_logging(level="DEBUG", debug=True)
     run_prediction_only(
         path_checkpoint=path_checkpoint,
         path_outputs=path_outputs,

src/segger/cli/segment.py

@@ -389,6 +389,7 @@ def segment(
 
     from lightning.pytorch.loggers import CSVLogger
     from lightning.pytorch import Trainer
+    from lightning.pytorch.callbacks import ModelCheckpoint
     from ..data import ISTSegmentationWriter
 
     csvlogger = CSVLogger(output_directory)
@@ -397,11 +398,23 @@ def segment(
         save_anndata=save_anndata,
         debug=debug,
     )
+    callbacks = [writer]
+
+    if debug:
+        checkpoint_callback = ModelCheckpoint(
+            dirpath=Path(output_directory) / "checkpoints",
+            filename="epoch={epoch:02d}",
+            save_top_k=-1,
+            every_n_epochs=1,
+        )
+        callbacks.append(checkpoint_callback)
+
+
     trainer = Trainer(
         logger=csvlogger,
         max_epochs=n_epochs,
         reload_dataloaders_every_n_epochs=1,
-        callbacks=[writer],
+        callbacks=callbacks,
     )
 
     # Training

src/segger/data/data_module.py

@@ -240,12 +240,18 @@ def load(self):
         )
 
         # Tile graph dataset
-        self.logger.debug("Tiling graph dataset...")
         node_positions = torch.vstack([
             self.data['tx']['pos'],
             self.data['bd']['pos'],
         ])
+        self.logger.debug(
+            f"Tiling graph: {len(node_positions)} positions, "
+            f"mode='{self.tiling_mode}'"
+        )
         if self.tiling_mode == "adaptive":
+            self.logger.debug(
+                f"  → QuadTreeTiling (max_tile_size={self.tiling_nodes_per_tile})"
+            )
             self.tiling = QuadTreeTiling(
                 positions=node_positions,
                 max_tile_size=self.tiling_nodes_per_tile,