diff --git a/src/anomalib/data/base/dataset.py b/src/anomalib/data/base/dataset.py
index 3f617ee004..581696983b 100644
--- a/src/anomalib/data/base/dataset.py
+++ b/src/anomalib/data/base/dataset.py
@@ -128,6 +128,7 @@ def __getitem__(self, index: int) -> dict[str, str | torch.Tensor]:
             # Therefore, create empty mask for Normal (0) images.
 
             mask = np.zeros(shape=image.shape[:2]) if label_index == 0 else cv2.imread(mask_path, flags=0) / 255.0
+            mask = mask.astype(np.single)
 
             transformed = self.transform(image=image, mask=mask)
 
diff --git a/src/anomalib/utils/cv/connected_components.py b/src/anomalib/utils/cv/connected_components.py
index 1f1f099dec..e2fc1000df 100644
--- a/src/anomalib/utils/cv/connected_components.py
+++ b/src/anomalib/utils/cv/connected_components.py
@@ -41,7 +41,7 @@ def connected_components_cpu(image: torch.Tensor) -> torch.Tensor:
     components = torch.zeros_like(image)
     label_idx = 1
     for i, msk in enumerate(image):
-        mask = msk.squeeze().numpy().astype(np.uint8)
+        mask = msk.squeeze().cpu().numpy().astype(np.uint8)
         _, comps = cv2.connectedComponents(mask)
         # remap component values to make sure every component has a unique value when outputs are concatenated
         for label in np.unique(comps)[1:]: