diff --git a/ultralytics/yolo/data/base.py b/ultralytics/yolo/data/base.py
index 1745f8f..d702b1b 100644
--- a/ultralytics/yolo/data/base.py
+++ b/ultralytics/yolo/data/base.py
@@ -110,7 +110,7 @@ class BaseDataset(Dataset):
                 if segments:
                     self.labels[i]["segments"] = segments[j]
             if self.single_cls:
-                self.labels[i]["cls"] = 0
+                self.labels[i]["cls"][:, 0] = 0
 
     def load_image(self, i):
         # Loads 1 image from dataset index 'i', returns (im, resized hw)
@@ -191,7 +191,7 @@ class BaseDataset(Dataset):
         return label
 
     def __len__(self):
-        return len(self.im_files)
+        return len(self.labels)
 
     def update_labels_info(self, label):
         """custom your label format here"""