diff --git a/utils/train_segmentation_utils.py b/utils/train_segmentation_utils.py
index e290c0e..c893044 100644
--- a/utils/train_segmentation_utils.py
+++ b/utils/train_segmentation_utils.py
@@ -76,8 +76,7 @@ def process_img_train(self, ind: float) -> np.array:
         :return: volume and labels (output mask)
         """
 
-        img_vert = self.array_training_img[ind]
-
+        img_vert = self.array_training_img[ind].astype(int)
         patch_nifti = nib.load(self.training_raw[img_vert[0]])
         mask_patch_nifti = nib.load(self.training_derivatives[img_vert[0]])