diff --git a/api/src/yolo/train_yolo_faces.py b/api/src/yolo/train_yolo_faces.py
index 63beb59..42d091d 100644
--- a/api/src/yolo/train_yolo_faces.py
+++ b/api/src/yolo/train_yolo_faces.py
@@ -34,7 +34,6 @@ def visualize_boxes(label, image, threshold=0.95):
     boxes_to_draw, grids_to_draw_obj, grids_to_draw_noobj = convert_prediction(label, image, threshold)
     if len(boxes_to_draw) == 0:
         print("no labels")
-        return
     boxes_to_draw = torch.tensor(boxes_to_draw)
     grids_to_draw_noobj = torch.tensor(grids_to_draw_noobj)
     grids_to_draw_obj = torch.tensor(grids_to_draw_obj)
diff --git a/api/src/yolo/yolo_dataset.py b/api/src/yolo/yolo_dataset.py
index b8a0d36..f2a0943 100644
--- a/api/src/yolo/yolo_dataset.py
+++ b/api/src/yolo/yolo_dataset.py
@@ -119,9 +119,6 @@ class YoloDataset(Dataset):
         boxes = augmented["bboxes"]
         labels = augmented["labels"]
 
-        if len(boxes) == 0:
-            return self.__getitem__((idx + 1) % len(self.image_ids))
-
         ground_truth = list(zip(boxes, labels))
         #[S, S, (x+y+w+h+c+C)]
         targets = np.zeros((self.grid, self.grid, 5 + self.num_classes), dtype=np.float32)