diff --git a/utils/dataset.py b/utils/dataset.py
index 45d31b7..bfba1a6 100644
--- a/utils/dataset.py
+++ b/utils/dataset.py
@@ -43,8 +43,8 @@ class BasicDataset(Dataset):
 
     def __getitem__(self, i):
         idx = self.ids[i]
-        mask_file = glob(self.masks_dir + idx + '*')
-        img_file = glob(self.imgs_dir + idx + '*')
+        mask_file = glob(self.masks_dir + idx + '.*')
+        img_file = glob(self.imgs_dir + idx + '.*')
 
         assert len(mask_file) == 1, \
             f'Either no mask or multiple masks found for the ID {idx}: {mask_file}'