Dice update

Author: C-Achard

napari_cellseg3d/code_models/worker_training.py

@@ -484,9 +484,7 @@ def get_loader_func(num_samples):
                 patience=self.config.scheduler_patience,
                 verbose=VERBOSE_SCHEDULER,
             )
-            dice_metric = DiceMetric(
-                include_background=False, reduction="mean"
-            )
+            dice_metric = DiceMetric(include_background=True, reduction="mean")
 
             best_metric = -1
             best_metric_epoch = -1
@@ -664,8 +662,8 @@ def get_loader_func(num_samples):
                         )
                         checkpoint_output = [
                             item.numpy()
-                            for batch in checkpoint_output
-                            for item in batch
+                            for channel in checkpoint_output
+                            for item in channel
                         ]
                         checkpoint_output[2] = checkpoint_output[2].astype(
                             np.uint16