[normalizer] Ensure targets are int type

src/mednet/libs/common/models/normalizer.py

@@ -3,12 +3,16 @@
 # SPDX-License-Identifier: GPL-3.0-or-later
 """Functions to compute normalisation factors based on dataloaders."""
 
+import logging
+
 import torch
 import torch.nn
 import torch.utils.data
 import torchvision.transforms
 import tqdm
 
+logger = logging.getLogger("mednet")
+
 
 def make_z_normalizer(
     dataloader: torch.utils.data.DataLoader,
@@ -31,6 +35,19 @@ def make_z_normalizer(
 
     # Peek the number of channels of batches in the data loader
     batch = next(iter(dataloader))
+
+    # Ensure targets are ints
+    try:
+        target = batch[1]["label"][0].item()
+        if not isinstance(target, int):
+            logger.info(
+                "Targets are not Integer type, skipping z-normalization."
+            )
+            return None
+    except RuntimeError:
+        logger.info("Targets are not Integer type, skipping z-normalization.")
+        return None
+
     channels = batch[0].shape[1]
 
     # Initialises accumulators